How to create a custom Dataset in PyTorch

Custom Dataset classes are the handoff between project-specific records and PyTorch training code. They are useful when samples need to be assembled from files, database rows, or precomputed feature arrays before a model can consume them.

In torch.utils.data.Dataset, getitem() returns one sample for an index, and len() gives samplers and DataLoader the dataset size. Keeping those methods small makes the class easier to batch, shuffle, split, and test before it reaches a training loop.

An in-memory smoke class keeps the dataset contract visible before any parser or storage layer is added. Replace the list-building layer with a CSV reader, image path lookup, or database query later, while keeping each returned sample shaped the same way.

1. Create custom_dataset.py with a map-style Dataset subclass and two smoke checks.

   custom_dataset.py

   from argparse import ArgumentParser
    
   import torch
   from torch.utils.data import DataLoader, Dataset
    
    
   class ChurnDataset(Dataset):
       def __init__(self, feature_rows, label_rows):
           if len(feature_rows) != len(label_rows):
               raise ValueError("feature_rows and label_rows must have the same length")
           self.features = torch.as_tensor(feature_rows, dtype=torch.float32)
           self.labels = torch.as_tensor(label_rows, dtype=torch.long)
    
       def __len__(self):
           return len(self.labels)
    
       def __getitem__(self, index):
           return {
               "features": self.features[index],
               "label": self.labels[index],
           }
    
    
   def build_dataset():
       feature_rows = [
           [22, 0, 0.10],
           [36, 1, 0.75],
           [41, 0, 0.35],
           [58, 1, 0.92],
       ]
       label_rows = [0, 1, 0, 1]
       return ChurnDataset(feature_rows, label_rows)
    
    
   def print_sample():
       dataset = build_dataset()
       sample = dataset[1]
       print(f"dataset_length={len(dataset)}")
       print(f"sample_keys={list(sample.keys())}")
       print(f"sample_features={sample['features'].tolist()}")
       print(f"sample_label={sample['label'].item()}")
    
    
   def print_batch():
       dataset = build_dataset()
       loader = DataLoader(dataset, batch_size=2, shuffle=False)
       batch = next(iter(loader))
       print(f"batch_features_shape={tuple(batch['features'].shape)}")
       print(f"batch_labels={batch['label'].tolist()}")
    
    
   def main():
       parser = ArgumentParser()
       parser.add_argument("--check", choices=["sample", "batch"], default="sample")
       args = parser.parse_args()
    
       if args.check == "sample":
           print_sample()
       else:
           print_batch()
    
    
   if __name__ == "__main__":
       main()

   The default DataLoader sampler sends integer indices to getitem(). Use a custom sampler only when the dataset key space is not integer-based.

2. Run the direct sample check.

   $ python custom_dataset.py --check sample
   dataset_length=4
   sample_keys=['features', 'label']
   sample_features=[36.0, 1.0, 0.75]
   sample_label=1

   dataset_length=4 comes from len(). The sample keys and tensor values come from getitem(1).

3. Run the DataLoader batch check.

   $ python custom_dataset.py --check batch
   batch_features_shape=(2, 3)
   batch_labels=[0, 1]

   batch_features_shape=(2, 3) shows two samples with three features each. batch_labels=[0, 1] confirms the loader collated labels from separate dataset items.
   Related: How to use a DataLoader in PyTorch