Module 4 — Dataset and DataLoader: feeding the training loop

The previous module built the Fashion-MNIST classifier as an nn.Module ready to consume batches of shape (batch, 1, 28, 28). The question this module answers is where those batches come from: how 60 000 files on disk, or a raw NumPy array in memory, become a stream of tensors that arrive on time, in the right shape, split cleanly between training and validation. Getting this part right is the difference between a training run that saturates the accelerator and one that spends 70 % of its time waiting on a disk.

Two objects, one contract

PyTorch splits data feeding in two.

Dataset describes one example at a time. Its whole interface is __len__ and __getitem__. Reading a specific image, applying a transform, returning the label: everything happens per example.

DataLoader describes how to assemble a batch. Given a dataset, it takes care of shuffling indices, drawing batches, spawning worker processes to prefetch, moving batches to pinned memory. Its interface is essentially "give me an iterator".

This separation is what makes PyTorch data loading composable: the same DataLoader code works with a torchvision dataset, a custom one over your own files, or a synthetic generator. Only Dataset changes; the training loop does not.

The built-in path: Fashion-MNIST in three lines

For classical datasets, torchvision gives you a Dataset already implemented.

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.ToTensor(),                    # PIL to (C, H, W) in [0, 1]
    transforms.Normalize((0.2860,), (0.3530,)),  # Fashion-MNIST stats
])

train_set = datasets.FashionMNIST("data", train=True, download=True, transform=transform)
val_set   = datasets.FashionMNIST("data", train=False, download=True, transform=transform)

train_loader = DataLoader(train_set, batch_size=64, shuffle=True, num_workers=2, pin_memory=True)
val_loader   = DataLoader(val_set,   batch_size=64, shuffle=False)

Three details already worth pausing on. ToTensor moves the layout to channels-first and rescales to $[0, 1]$: the network from module 3 expects exactly that. Normalize subtracts the mean and divides by the standard deviation, computed once on the training set. And shuffle=True is on for training, off for validation — an evaluation with shuffled batches produces the same numbers but makes any per-batch diagnostic impossible.

The normalisation leak: the number one data mistake

The most dangerous line in the block above is the one that looks most innocent: the tuple (0.2860,), (0.3530,).

Those numbers are the mean and standard deviation of the training set. Applying them to the validation set is correct. Computing them on the full dataset before splitting is a leak — the validation set has contributed to a preprocessing step that will run at inference time, and the reported validation accuracy is optimistic in a way that only shows up in production.

The rule is unambiguous: fit any statistic on the training set alone, then apply it to validation, test and every future example. This includes means, standard deviations, tokenisers, PCA rotations, TF-IDF vocabularies, and every scaler ever imagined. The datasets.FashionMNIST shortcut hides this because Fashion-MNIST comes pre-split; on your own data, the split must come first.

This is the failure mode reviewers spot first

When a model that reported 94 % on validation degrades to 87 % in production, the first hypothesis is a normalisation leak. Auditing preprocessing is the very first sanity check, before hyperparameters, before architecture, before anything else. The fix is architectural: split first, then compute statistics inside a training-only pipeline.

A custom Dataset in twenty lines

When the data is not standard, you write your own Dataset. This is much simpler than it sounds.

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

class FashionFolder(Dataset):
    def __init__(self, root, transform=None):
        self.paths = sorted(Path(root).glob("*/*.png"))
        self.classes = sorted({p.parent.name for p in self.paths})
        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}
        self.transform = transform

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

    def __getitem__(self, idx):
        path = self.paths[idx]
        image = Image.open(path).convert("L")
        label = self.class_to_idx[path.parent.name]
        if self.transform is not None:
            image = self.transform(image)
        return image, label

Three points to remember. __init__ should only index the data — collect paths, read labels, precompute what is cheap — never load the actual pixels: __init__ runs on the main process, and loading 60 000 images there would delay the first batch by minutes. __getitem__ does the heavy lifting per example, and it is that method that is called in worker processes. Anything expensive belongs inside __getitem__, not outside.

DataLoader arguments that matter

Argument	What it does	When it changes
batch_size	number of examples per batch	fits the accelerator memory
shuffle	permutes indices each epoch	True for training, False for validation
num_workers	subprocesses that call __getitem__	as many as CPU cores minus one
pin_memory	allocates page-locked memory	True for GPU training, ignored otherwise
drop_last	drops the trailing partial batch	True for BatchNorm-heavy models

num_workers=0 means loading runs in the main process, and every batch pauses training while it is prepared. Any positive value spawns subprocesses that prefetch batches in parallel with the accelerator's work — this is exactly the pipeline pattern from course 08, applied to PyTorch. The optimal count is machine-specific; two to eight is the usual range on a laptop.

pin_memory=True allocates batches in page-locked host memory, which makes the subsequent .to("cuda", non_blocking=True) transfer overlap with computation. Without it, the transfer stalls the GPU. It is the second most impactful DataLoader setting after num_workers.

Train, validation and test: three sets, one code path

The most defensible split is created once, saved, and reused across every experiment. A three-way split for Fashion-MNIST could look like:

from torch.utils.data import random_split

torch.manual_seed(42)
n = len(train_set)
train_data, val_data = random_split(train_set, [55_000, 5_000])
test_data = val_set     # the held-out FashionMNIST test split

Test is held out entirely. Validation drives every decision — learning rate, architecture, when to stop. Reporting the test number happens once, at the end, and never guides a hyperparameter search. A team that "tries a few things on the test set" is running validation with extra steps and reporting numbers that will not survive first contact with new data.

collate_fn: assembling the batch

The default collate_fn stacks tensors along a new leading dimension and turns lists of labels into a LongTensor. That covers 90 % of cases. You override it when examples have variable lengths — text sequences, variable-size images — and you need to pad or bucket them. In this course we stick to the default; module 4 of the NLP course revisits collate_fn for text.

Time one epoch with num_workers=0 and again with num_workers=4

If the second is two or three times faster, the data loading is your bottleneck and this module's advice matters. If they are close, the accelerator is saturated by compute and workers are not the problem. Measuring before optimising remains the golden rule.

In summary

• Dataset returns one example, DataLoader returns a batch; the split is what makes PyTorch data loading composable across projects.
• The normalisation leak — computing statistics on train and validation together — is the single most damaging preprocessing mistake, because it produces optimistic validation numbers that collapse in production.
• num_workers > 0 and pin_memory=True are the two settings that let training and data preparation overlap; without them, the accelerator waits.
• Test is held out, validation drives every decision, and the split is created once and reused across experiments to keep reported numbers comparable.

Next module: the training loop itself, which finally consumes these batches and turns the Fashion-MNIST classifier into a network that actually learns.