How to create a custom Dataset in PyTorch

Training code stays easier to test when sample lookup is separated from model logic. A PyTorch Dataset provides that boundary for project-specific records while preserving the indexed access expected by the training input pipeline.

A map-style Dataset returns one sample from getitem() and reports its size through len(). The default DataLoader sampler supplies integer indexes, then its collate function combines matching fields from the returned samples.

The in-memory records below keep storage and parsing out of the way so the class contract remains visible. Each item returns a dictionary with a three-value floating-point feature tensor and one integer class label, which lets the default collate function create batched tensors without a custom handler.

1. Create custom_dataset.py with the ChurnDataset initializer.

   custom_dataset.py

   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)

2. Add the len() method below the initializer.

       def __len__(self):
           return len(self.labels)

3. Add the getitem() method below len().

       def __getitem__(self, index):
           return {
               "features": self.features[index],
               "label": self.labels[index],
           }

4. Append the sample records with a direct item check after the class.

   feature_rows = [
       [22, 0, 0.10],
       [36, 1, 0.75],
       [41, 0, 0.35],
       [58, 1, 0.92],
   ]
   label_rows = [0, 1, 0, 1]
    
   dataset = ChurnDataset(feature_rows, label_rows)
   sample = dataset[1]
    
   assert len(dataset) == 4
   assert sample["features"].shape == torch.Size([3])
   assert sample["label"].dtype == torch.long
    
   print(f"dataset_length={len(dataset)}")
   print(f"sample_features_shape={tuple(sample['features'].shape)}")
   print(f"sample_label_dtype={sample['label'].dtype}")

5. Append the DataLoader batch check below the direct item check.

   loader = DataLoader(dataset, batch_size=2, shuffle=False)
   batch = next(iter(loader))
    
   assert batch["features"].shape == torch.Size([2, 3])
   assert batch["label"].shape == torch.Size([2])
    
   print(f"batch_features_shape={tuple(batch['features'].shape)}")
   print(f"batch_labels={batch['label'].tolist()}")

   batch_size=2 asks the default collate function to stack two returned dictionaries. shuffle=False keeps the first batch in source order.
   Related: How to use a DataLoader in PyTorch

6. Run custom_dataset.py to verify indexed samples collate into a two-row batch.

   $ python3 custom_dataset.py
   dataset_length=4
   sample_features_shape=(3,)
   sample_label_dtype=torch.int64
   batch_features_shape=(2, 3)
   batch_labels=[0, 1]