PyTorch Custom Dataset

Built-in datasets like MNIST and CIFAR-10 work for learning, but real projects require loading your own data — CSVs, images in custom folder structures, audio files, or records from a database. Writing a custom Dataset class gives you full control over how data is read, preprocessed, and served to the training loop.

Custom Dataset Template

import torch
from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, ...):
        # Load or reference data here
        # Do NOT load individual samples yet — just store paths or arrays
        pass

    def __len__(self):
        # Return total number of samples
        return ...

    def __getitem__(self, idx):
        # Load and return the sample at position idx
        # Apply transforms here if needed
        return sample, label

Custom Dataset from a CSV File

A CSV file with feature columns and a label column is one of the most common data sources in practice.

import torch
import pandas as pd
from torch.utils.data import Dataset
from sklearn.preprocessing import StandardScaler

class CSVDataset(Dataset):
    def __init__(self, filepath, label_col='target'):
        df = pd.read_csv(filepath)
        self.labels = torch.tensor(df[label_col].values, dtype=torch.long)

        features = df.drop(columns=[label_col]).values.astype('float32')
        scaler   = StandardScaler()
        self.features = torch.tensor(scaler.fit_transform(features))

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

    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]

dataset = CSVDataset('data.csv')
print(f"Samples: {len(dataset)}")
print(f"Sample shape: {dataset[0][0].shape}")

Custom Image Dataset

This example loads images from a folder where each subfolder is a class. It reads image paths at init time and loads image pixels only when __getitem__ is called — this is called lazy loading and avoids loading the entire dataset into RAM.

import os
from PIL import Image
import torch
from torch.utils.data import Dataset
from torchvision import transforms

class ImageDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.transform  = transform
        self.samples    = []   # list of (image_path, label_index)
        self.classes    = sorted(os.listdir(root_dir))
        self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)}

        for cls in self.classes:
            cls_dir = os.path.join(root_dir, cls)
            for fname in os.listdir(cls_dir):
                if fname.lower().endswith(('.jpg', '.png', '.jpeg')):
                    path = os.path.join(cls_dir, fname)
                    self.samples.append((path, self.class_to_idx[cls]))

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

    def __getitem__(self, idx):
        path, label = self.samples[idx]
        img = Image.open(path).convert('RGB')
        if self.transform:
            img = self.transform(img)
        return img, label

transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor()
])

dataset = ImageDataset(root_dir='./my_images', transform=transform)
print(f"Classes: {dataset.classes}")
print(f"Total images: {len(dataset)}")

Dataset Folder Structure This Expects

my_images/
├── cats/
│   ├── img001.jpg
│   └── img002.jpg
├── dogs/
│   ├── img003.jpg
│   └── img004.jpg
└── birds/
    └── img005.jpg

class_to_idx: {'birds': 0, 'cats': 1, 'dogs': 2}

Lazy vs Eager Loading

Eager Loading (load all data in __init__):
  ✓ Fast access per sample (already in RAM)
  ✗ High RAM usage — fails for large datasets

Lazy Loading (load sample in __getitem__):
  ✓ Works for any dataset size — no RAM limit
  ✗ Disk I/O per batch — use num_workers to compensate

Rule: Use lazy loading for image/audio/large datasets.
      Use eager loading only when dataset fits in RAM.

Adding Transforms Conditionally

A common pattern is to pass different transforms at init time depending on whether the dataset instance is for training or validation.

from torchvision import transforms

train_transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.5]*3, [0.5]*3)
])

val_transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
    transforms.Normalize([0.5]*3, [0.5]*3)
])

train_ds = ImageDataset('./data/train', transform=train_transform)
val_ds   = ImageDataset('./data/val',   transform=val_transform)

Using Your Custom Dataset with DataLoader

from torch.utils.data import DataLoader

train_loader = DataLoader(train_ds, batch_size=32, shuffle=True,  num_workers=2)
val_loader   = DataLoader(val_ds,   batch_size=32, shuffle=False, num_workers=2)

for images, labels in train_loader:
    print(images.shape)   # torch.Size([32, 3, 128, 128])
    print(labels.shape)   # torch.Size([32])
    break

Caching Preprocessed Data

If preprocessing is slow (e.g. loading and resizing thousands of high-resolution images), cache the preprocessed tensors to disk after the first run and load the cache on subsequent runs. This trades disk space for much faster training startup.

import os

class CachedDataset(Dataset):
    def __init__(self, source_dataset, cache_path):
        self.cache_path = cache_path
        if os.path.exists(cache_path):
            self.data = torch.load(cache_path)
        else:
            self.data = [source_dataset[i] for i in range(len(source_dataset))]
            torch.save(self.data, cache_path)

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

    def __getitem__(self, idx):
        return self.data[idx]

Summary

A custom Dataset subclass requires __len__ and __getitem__. Store file paths or arrays in __init__ and load individual samples lazily in __getitem__ to handle large datasets without filling RAM. Apply transforms inside __getitem__ for on-the-fly processing. Pass your custom Dataset to a DataLoader exactly like any built-in dataset. Use separate Dataset instances (or transform arguments) for training and validation so augmentation applies only during training. Custom Datasets give you complete control over any data source — file formats, preprocessing steps, label encoding, and caching.