pytorch Subset to return an instance of the parent Dataset, to be able to access the same attribute

torch.utils.data.dataset.py

"""
An implementation of the pytorch Subset that returns an instance of the original dataset with a reduced number of items.
This has two benefits:
  - It allows to stil access the attributes of the Dataset class, such as methods, or properties.
  - You can use the usual python index notation with slices to chunk the dataset, rather than creating a list of indices

"""
class Dataset(object):
  def __init__(self, iterable):
    self.items = iterable
    self.subset = Subset(self)
    
    self.any_field = NotImplemented
    
  def __getitem__(self, idx):
    return self.items[idx]  # or whatever you want
    
  def __len__(self):
    return len(self.items)  # or whatever you want
  
  def any_function(self, *args, **kwargs):
    raise NotImplementedError
    
class Subset(object):
    def __init__(self, dataset):
        self.dataset = dataset

    def __getitem__(self, idx):
        subset = self.dataset
        subset.items = subset.items[idx]
        return subset

# usage:
if __name__ == "__main__":
  dataset = Dataset(list(range(100)))
  subset = dataset.subset[0:10]
  # you can now still access <any_field> and <any_function> from subset
  # where the usual pytorch Subset implementation would have raised AttributeError
  any_field = subset.any_field
  any_function = subset.any_function