3. Dataset and DataLoader engineering

In the previous chapter we established that a PyTorch training loop alternates between forward passes, backward passes, and optimizer steps. What often goes unnoticed is that the loop can only run as fast as its slowest component. On modern GPUs, a single forward-backward step on a ResNet-50 with batch size 256 may take 30–50 ms, while naively loading and decoding 256 JPEG images on the CPU can easily take 200 ms or more. The result is a GPU that sits idle most of the time, waiting for data. This chapter focuses on the engineering discipline of feeding the model fast enough: how Dataset and DataLoader interact, how transforms and samplers shape the data stream, and how prefetching, worker processes, and memory pinning keep the accelerator saturated.

The Dataset Contract

A PyTorch Dataset is any object that implements two methods: __len__, returning the number of samples, and __getitem__(index), returning one sample (or a tuple of tensors). The map-style interface is the most common, but PyTorch also supports IterableDataset for streaming sources such as log files or remote object stores where random access is impossible.

A minimal map-style dataset for image classification looks like this:

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

class ImageFolderDataset(Dataset):
    def __init__(self, samples, transform=None):
        # samples: list of (image_path, label)
        self.samples = samples
        self.transform = transform

    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

The critical engineering point is that __getitem__ runs inside worker processes when num_workers > 0. Any state it touches must be picklable, and any file handle or database connection should be opened lazily inside the method rather than in __init__. A common bug is opening a file in __init__ and then forking: the child processes inherit a stale descriptor and read corrupted data.

Transforms: Where CPU Time Goes

Transforms are the dominant CPU cost in most vision pipelines. Decoding a JPEG, resizing, and normalizing is far more expensive than the tensor copy that follows. Two rules of thumb apply:

• Do the cheap, tensor-friendly operations (normalization, channel permutation) after conversion to tensor; do the expensive geometric operations (resize, crop, flip) on the PIL or NumPy representation, which is faster for those kernels.
• Use torchvision.transforms.v2 when possible. It operates on tensors directly and fuses several operations, reducing Python overhead.

For training-time augmentation, the transform pipeline is stochastic. For validation, it must be deterministic. A typical configuration:

from torchvision import transforms as T

train_tf = T.Compose([
    T.RandomResizedCrop(224, scale=(0.6, 1.0)),
    T.RandomHorizontalFlip(),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]),
])

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

Note that the normalization constants are ImageNet statistics. Using the wrong mean and std will not crash training, but it will slow convergence and can cost several points of accuracy. Always verify that the statistics match the pretrained weights you load.

Samplers: Controlling What the Model Sees

The DataLoader does not decide the order of samples; the sampler does. By default, DataLoader uses RandomSampler when shuffle=True and SequentialSampler otherwise. For imbalanced datasets, WeightedRandomSampler lets you oversample rare classes:

from torch.utils.data import WeightedRandomSampler

class_counts = [900, 80, 20]          # e.g., three classes
weights = [1.0 / c for c in class_counts]
sample_weights = [weights[label] for _, label in samples]

sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=len(sample_weights),
    replacement=True,
)

For distributed training, DistributedSampler shards the dataset across ranks so that each GPU sees a disjoint subset. It must be told the current epoch via sampler.set_epoch(epoch); otherwise every epoch shuffles identically and the shards never rotate. This is one of the most common silent bugs in multi-GPU training.

Prefetching, Workers, and Pinned Memory

The DataLoader overlaps data loading with GPU computation through three mechanisms:

1. Worker processes (num_workers): each worker runs __getitem__ independently, so decoding is parallelized across CPU cores.
2. Prefetching (prefetch_factor): each worker prepares prefetch_factor batches ahead of time, so the main process rarely blocks.
3. Pinned memory (pin_memory=True): the batch is copied into page-locked host memory, enabling asynchronous host-to-device transfers via tensor.to(device, non_blocking=True).

A practical starting configuration for a single-GPU workstation with 8 CPU cores:

from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=128,
    shuffle=True,
    num_workers=4,
    prefetch_factor=2,
    pin_memory=True,
    persistent_workers=True,
    drop_last=True,
)

persistent_workers=True keeps worker processes alive between epochs, avoiding the fork/join cost that can dominate when epochs are short. drop_last=True discards the final partial batch, which prevents BatchNorm from seeing a batch of size 1 and producing unstable statistics.

Putting It Together: A Data Flow View

The diagram below traces a single batch from disk to GPU. Each stage has a distinct cost profile, and the engineering goal is to keep the GPU stage from ever waiting on the earlier ones.

Diagnosing and Fixing Common Pipeline Failures

Three failure modes account for most slow or broken input pipelines:

Scaling Beyond a Single Machine

When the dataset no longer fits on local disk, the pipeline must be re-architected. Two patterns dominate:

• Sharded webdataset: pack samples into large tar shards (e.g., 1 GB each) and stream them with an IterableDataset. This avoids the per-file open/close overhead that kills performance on object stores.
• Pre-tokenized / pre-decoded caches: for NLP, tokenize once and store as memory-mapped arrays; for vision, store decoded tensors in a WebDataset or LMDB. The transform pipeline then only does augmentation, not decoding.

In both cases, the sampler must be replaced by a sharding strategy: each rank reads a disjoint subset of shards, and the IterableDataset yields samples until exhausted. The DistributedSampler does not apply to iterable datasets, so you must implement shard assignment manually, typically by passing rank and world_size to the dataset constructor.

Key Takeaways

• The DataLoader is a pipeline, not a convenience wrapper. Its throughput is determined by the slowest stage: I/O, decode, transform, collate, or transfer.
• Set num_workers to roughly the number of physical CPU cores, prefetch_factor to 2–4, and pin_memory=True with non_blocking=True for GPU training.
• Use persistent_workers=True and drop_last=True to avoid epoch-boundary stalls and BatchNorm instability.
• Choose the sampler deliberately: WeightedRandomSampler for class imbalance, DistributedSampler for multi-GPU, and manual sharding for IterableDataset.
• Profile before optimizing. If GPU utilization is already above 85%, the data pipeline is not your bottleneck.