6.9 The training loop, datasets, dataloaders

What this is and why it exists

The training loop is the instrument you will use for every project in this module and afterwards, so build one good version and reuse it. Written well, a new project becomes filling in blanks: swap the dataset, swap the model, change a few settings. Written badly, every project re-derives the same mistakes. This topic also covers the failure that wastes the most money in practice — an accelerator sitting idle because the data could not arrive fast enough, blamed on the model for weeks.

The vocabulary

• Dataset — an object that knows how to produce one example.
• DataLoader — the object that batches, shuffles and parallelises the fetching.
• Collate function — the function that assembles individual examples into a batch.
• Epoch — one pass over the training data.
• Checkpoint — a saved snapshot from which training can continue.
• Seed — the number that makes a random sequence repeatable.
• Determinism — whether the same inputs give bit-identical outputs.
• Throughput — examples processed per second, end to end.

The mental model

Two objects stand between raw data and a batch. A dataset produces one example; a loader turns examples into batches. PyTorch distinguishes two dataset styles: map-style datasets implement item access and length and "represent a map from (possibly non-integral) indices/keys to data samples", while iterable-style datasets are "particularly suitable for cases where random reads are expensive or even improbable" — a distinction that decides itself, since files on disk are random-access and a network stream is not.

Keep the dataset's job narrow: read one item, apply the transformations, return it. Reading a whole file per item, or holding a database connection created before the worker processes fork, is where dataset code goes wrong.

The collate function is needed more often than beginners expect. The default stacks examples that are all the same shape, which covers fixed-size images and nothing else. Variable-length sequences must be padded to the longest in the batch and accompanied by a mask saying which positions are real; detection targets differ in count per image and cannot be stacked at all. Writing a custom one is routine, and the moment you meet a shape error that mentions stacking, this is what it is asking for.

The loader's settings are where throughput is won or lost. The worker count is documented as "how many subprocesses to use for data loading", with zero meaning the main process — which means the default loads data in the same process that is meant to be feeding the accelerator, and the accelerator waits. Pinned memory, which makes the loader "copy Tensors into device/CUDA pinned memory before returning them", speeds the transfer. Shuffling is documented as reshuffling "at every epoch" and must be on for training and off for validation, so a validation number is comparable between epochs. Dropping the last incomplete batch matters where a batch of one would break a normalisation layer. And persistent workers, which keep the processes alive between epochs, remove a startup cost that is significant when epochs are short.

The loop itself has a fixed shape, and writing it once properly is the point of the topic.

for epoch in range(epochs):
    model.train()
    for batch in train_loader:
        loss = criterion(model(batch.x), batch.y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    model.eval()
    with torch.no_grad():
        for batch in val_loader:
            ...  # accumulate metrics as plain numbers, not live tensors

Four things belong around that skeleton. Explicit mode switching, because of the normalisation behaviour from the earlier topic. No-grad in validation, for the memory reason. Metrics recorded as plain numbers, because accumulating live tensors keeps every batch's graph alive. And logging that includes throughput and the learning rate alongside the loss, because those two numbers answer most questions about why a run is behaving oddly.

Checkpointing is what makes long runs survivable, and the rule is that a checkpoint must contain everything needed to continue rather than everything needed to predict. That means the model's parameters, the optimiser's state, the scheduler's position, the epoch number, and the configuration. Resume without the optimiser state and the momentum and adaptive estimates restart from nothing; resume without the scheduler's position and the learning rate jumps back to where it was at the start. Save the most recent checkpoint and separately the best one by validation score, and write to a temporary file before renaming, so an interrupted save does not leave you with a corrupt file where your last good one was.

Reproducibility is a real trade rather than an ideal. Seeding the random number generators makes a run repeatable in the ordinary sense — same initialisation, same shuffling, same augmentation. Full determinism is more than that: some accelerator operations accumulate in a nondeterministic order, and forcing deterministic alternatives is slower and, for a few operations, impossible. Decide what you actually need. For comparing two configurations, seeding is enough and running each seed three times tells you more than making one run bit-identical. For a result somebody must reproduce exactly, pay for determinism deliberately and record the versions, because determinism does not survive a library upgrade. And seed the loader's workers as well, or your augmentation is repeatable in the main process and random in the four that actually do the work.

Finally, the starvation failure. If the accelerator is idle half the time, doubling the model's speed buys you nothing, and no amount of architectural cleverness will show up in wall-clock time. Log examples per second from the first run, and watch accelerator utilisation. Low utilisation with high processor usage means the input pipeline is the bottleneck: add workers, move augmentation off the critical path, pre-resize images rather than resizing every epoch, or store the data in a format that reads in one operation instead of thousands. Measure the pipeline before optimising the model, because the input side is where the idle time usually is.

What you should now be able to explain or do

Distinguish the two dataset styles and say which suits random reads. Keep dataset code to reading one example. Write a collate function for variable-length or variably-shaped data. Set worker count, pinning, shuffling and last-batch behaviour deliberately, and say what each buys. Write a reusable loop with mode switching, no-grad validation, plain-number metrics and throughput logging. Save a checkpoint that can be resumed from, including optimiser and scheduler state, safely. Decide how much reproducibility you need and pay for that much. Diagnose data starvation from utilisation rather than guessing.

Check yourself