PyTorch Custom Datasets

This note builds a complete, realistic custom Dataset — loading images from a folder structure with labels from a CSV file — the kind of practical data-loading code every real project eventually needs.

A Realistic Example: Images From Disk + a Labels CSV

import os
import pandas as pd
from PIL import Image
from torch.utils.data import Dataset

class ImageLabelDataset(Dataset):
    def __init__(self, csv_path, image_dir, transform=None):
        self.labels_df = pd.read_csv(csv_path)   # columns: 'filename', 'label'
        self.image_dir = image_dir
        self.transform = transform

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

    def __getitem__(self, idx):
        row = self.labels_df.iloc[idx]
        image_path = os.path.join(self.image_dir, row['filename'])
        image = Image.open(image_path).convert('RGB')   # loaded LAZILY, on demand
        label = row['label']

        if self.transform:
            image = self.transform(image)

        return image, label

Usage With Transforms and a DataLoader

import torchvision.transforms as transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

dataset = ImageLabelDataset("labels.csv", "images/", transform=transform)
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

for images, labels in loader:
    print(images.shape)   # (32, 3, 224, 224) -- a full, ready-to-train batch
    break

Handling a Train/Val/Test Split

from torch.utils.data import random_split

full_dataset = ImageLabelDataset("labels.csv", "images/", transform=transform)
train_size = int(0.7 * len(full_dataset))
val_size = int(0.15 * len(full_dataset))
test_size = len(full_dataset) - train_size - val_size

train_ds, val_ds, test_ds = random_split(full_dataset, [train_size, val_size, test_size])

Common Mistakes

• Applying data augmentation transforms to validation/test datasets — augmentation (random crops, flips) belongs only on the training split; validation/test should use only deterministic preprocessing (resize, normalize) for consistent, comparable evaluation.
• Loading all images into memory inside __init__ for a large dataset — this defeats the purpose of lazy loading and can exhaust available memory; load each image only when __getitem__ is actually called for it.

Interview Relevance

Q: "Why should image loading happen inside __getitem__ rather than upfront in __init__ for a large image dataset?" Loading every image into memory upfront doesn't scale to large datasets — a dataset with millions of images would exhaust available RAM long before training even starts. Loading each image lazily, only when __getitem__ is called for that specific index, keeps memory usage proportional to the batch size rather than the entire dataset, and combined with DataLoader's num_workers, allows parallel loading to overlap with GPU computation.

Practice Question

Why must train-set augmentation transforms (random crop, flip) be excluded from the validation and test dataset's transform pipeline?

Related DL Notes

• PyTorch Dataset
• Data Augmentation
• Dataset Train/Val/Test Split