Data pipelines and augmentation with torchvision

Transforms and Compose, custom Datasets, weighted samplers for imbalanced classes, collate functions, and worker settings that do not stall.

A custom Dataset and the transform split

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

class ImageDataset(Dataset):
    def __init__(self, rows, transform=None):
        self.rows = rows                # list of (path, label)
        self.transform = transform

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

    def __getitem__(self, index):
        path, label = self.rows[index]
        image = Image.open(path).convert("RGB")
        if self.transform is not None:
            image = self.transform(image)
        return image, label

# augmentation only on the training split
train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.2, 0.2, 0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

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

• __getitem__ must be independent per index. Any shared mutable state across calls breaks with multiple workers, because each worker is a separate process.
• Forward slashes throughout. On Windows, a path exactly 260 characters long fails inside a worker with a bare FileNotFoundError; the fix is a shorter root or the long-path manifest.
• The normalisation constants above are the ImageNet statistics. Reusing them for a very different domain (medical, satellite) is a reasonable default but not an optimal one.
• ToTensor() converts an HWC uint8 image to a CHW float tensor scaled to [0, 1]. Anything after it operates on tensors, not PIL images.

Samplers, collate and DataLoader

from torch.utils.data import DataLoader, WeightedRandomSampler
import numpy as np

labels = np.array([row[1] for row in train_rows])
class_counts = np.bincount(labels)
class_weights = 1.0 / np.maximum(class_counts, 1)
sample_weights = class_weights[labels]                  # one weight per example

sampler = WeightedRandomSampler(
    weights=torch.as_tensor(sample_weights, dtype=torch.double),
    num_samples=len(sample_weights),
    replacement=True,
)

train_loader = DataLoader(
    ImageDataset(train_rows, train_tf),
    batch_size=32,
    sampler=sampler,            # mutually exclusive with shuffle=True
    num_workers=4,
    pin_memory=True,            # faster host-to-device copies on CUDA
    persistent_workers=True,    # do not restart workers each epoch
    drop_last=True,
    prefetch_factor=2,
)

# a collate function for variable-length sequences
def pad_collate(batch):
    sequences, labels = zip(*batch)
    lengths = torch.tensor([len(s) for s in sequences])
    padded = torch.nn.utils.rnn.pad_sequence(sequences, batch_first=True)
    mask = torch.arange(padded.size(1))[None, :] < lengths[:, None]
    return padded, lengths, mask, torch.stack(labels)

A WeightedRandomSampler balances what the model sees without touching the data, but it also changes the effective prior. Combine it with a class-weighted loss only if you have checked that the two effects do not overshoot.

Making the loader fast

import time

def benchmark(loader, n=30):
    it = iter(loader)
    next(it)                                  # warm up workers
    start = time.perf_counter()
    for _ in range(n):
        batch = next(it)
    elapsed = time.perf_counter() - start
    return elapsed / n, batch[0].shape

per_batch, shape = benchmark(train_loader)
print(f"{per_batch * 1000:.1f} ms per batch, last batch {tuple(shape)}")

# worker crash tips: a traceback from inside a worker is truncated to one frame.
# Set the environment before importing torch to get the real error.
import os
os.environ["PYTHONFAULTHANDLER"] = "1"
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"     # forces synchronous CUDA errors

• Compare batch time against model step time. If loading dominates, add workers; if not, look at the model or at host-to-device transfer.
• Pin a thread count when using many workers: torch.set_num_threads(1) inside workers avoids thread oversubscription and often doubles throughput.
• A worker may only be slow to fault. Errors inside a worker surface with a truncated traceback, so run once with num_workers=0 to get the real exception.
• persistent_workers=True avoids paying process startup every epoch, but it holds memory: on a machine running several jobs, that matters.

FAQ

Related

Modules, data and saving Mixed precision and GPU performance

Last refreshed 2026-09-18.