Data pipelines: Dataset, DataLoader & augmentation

Real projects rarely come as a neat torchvision.datasets download. You will have a folder of product photos, a CSV of transactions, or a database table. PyTorch's Dataset and DataLoader classes turn any of these into fast, shuffled batches — and a good pipeline is often the difference between a GPU that is 30% busy and one that is 95% busy.

The Dataset contract

A map-style dataset needs just two methods: __len__ (how many examples) and __getitem__ (return one example by index).

import pandas as pd
import torch
from torch.utils.data import Dataset

class DemandDataset(Dataset):
    # Each row: store_id, day_of_week, is_festival, temp_c, units_sold
    def __init__(self, csv_path: str):
        df = pd.read_csv(csv_path)
        self.x = torch.tensor(
            df[["day_of_week", "is_festival", "temp_c"]].to_numpy(),
            dtype=torch.float32,
        )
        self.y = torch.tensor(df["units_sold"].to_numpy(), dtype=torch.float32)

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

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

For small tabular data, converting everything to tensors once in __init__ is fastest. For images, load lazily in __getitem__ so you never hold thousands of photos in RAM.

Folders of images

If your images are arranged as root/class_name/image.jpg, ImageFolder does the work:

sweets/
  train/
    jalebi/  001.jpg 002.jpg ...
    kaju_katli/ ...
    rasgulla/ ...
  val/
    jalebi/ ...

from torchvision import datasets, transforms

train_tfms = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
val_tfms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

train_ds = datasets.ImageFolder("sweets/train", transform=train_tfms)
val_ds = datasets.ImageFolder("sweets/val", transform=val_tfms)
print(train_ds.class_to_idx)   # {'jalebi': 0, 'kaju_katli': 1, 'rasgulla': 2}

The normalisation values are the ImageNet channel means and standard deviations, required when using ImageNet-pretrained models.

A custom image dataset from a CSV

When labels live in a spreadsheet instead of folder names:

from pathlib import Path
from PIL import Image

class LabelledImages(Dataset):
    def __init__(self, csv_path, img_dir, transform=None):
        self.df = pd.read_csv(csv_path)          # columns: filename, label
        self.img_dir = Path(img_dir)
        self.classes = sorted(self.df["label"].unique())
        self.to_idx = {c: i for i, c in enumerate(self.classes)}
        self.transform = transform

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

    def __getitem__(self, i):
        row = self.df.iloc[i]
        img = Image.open(self.img_dir / row["filename"]).convert("RGB")
        if self.transform:
            img = self.transform(img)
        return img, self.to_idx[row["label"]]

The .convert("RGB") call protects you from grayscale or PNG-with-alpha images that would otherwise break batching.

DataLoader settings that matter

from torch.utils.data import DataLoader

train_dl = DataLoader(train_ds, batch_size=32, shuffle=True,
                      num_workers=4, pin_memory=True, drop_last=True)
val_dl = DataLoader(val_ds, batch_size=64, shuffle=False, num_workers=4)

Argument	Effect
batch_size	Larger is faster per epoch but uses more memory
shuffle	True for training only
num_workers	Parallel loading processes; try the number of CPU cores
pin_memory	Faster CPU-to-GPU copies
drop_last	Avoids a tiny last batch (useful with batch norm)

On Windows and macOS, code that uses num_workers > 0 must run under if __name__ == "__main__":.

Augmentation principles

Augmentation creates plausible variations so the model learns what matters. Choose transforms that preserve the label:

• Horizontal flip is fine for sweets or clothes, wrong for digits like 6 or text.
• Strong colour jitter can be harmful when colour is the signal (ripe vs unripe mango).
• Apply augmentation to training data only.

Handling class imbalance

If you have 2,000 kaju katli photos but only 150 of rasmalai, use a WeightedRandomSampler so each class is sampled roughly equally:

from collections import Counter
from torch.utils.data import WeightedRandomSampler

counts = Counter(train_ds.targets)
weights = [1.0 / counts[t] for t in train_ds.targets]
sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)
train_dl = DataLoader(train_ds, batch_size=32, sampler=sampler)   # no shuffle with sampler

Common mistakes

• Splitting data after augmentation, leaking near-duplicates into validation.
• Forgetting Normalize with pretrained models.
• Random splits for time-series data — always split by date instead.
• Near-identical photos (burst shots) in both train and validation, inflating accuracy.