Batching and the training loop

Epochs, steps and batch size

An epoch is one pass over the training set. A step is one update of the weights. A batch is the group of examples used for one update. With N examples and a batch of size B, an epoch is about N divided by B steps, rounded up. The last batch is often smaller unless you drop it. That one batch can make the loss curve look noisy at regular intervals.

The batch size trades noise for throughput. A small batch gives a noisy gradient that can escape shallow minima, and it fits in memory. A large batch gives a stable gradient and uses the hardware better. It may also need a higher learning rate and more epochs to reach the same place. Neither is universally better, and the choice is usually set by memory first.

The last batch is worth a thought. With drop_last=False it is smaller, so its gradient is noisier and its loss has a different denominator. With drop_last=True a few examples are never seen in that epoch. Neither is fatal. A metric averaged over batches without weighting will disagree with a metric averaged over examples, though.

The loop, line by line

The order is fixed. Clear the gradients, run the forward pass, compute the loss, run backward, then update the weights. Clearing before backward matters because gradients accumulate. Updating after backward matters because the optimiser reads the gradients that backward wrote a moment ago. Swap two lines and the run still executes.

Shuffle the data at the start of every epoch. Without it the model sees the same batches in the same order. The gradient noise then becomes periodic, which can look like a learning rate problem. A seeded shuffle keeps the run reproducible without keeping the order fixed.

The optimiser step is the only line that changes the weights. Everything before it is bookkeeping. A missing optimizer.step() therefore produces a loss curve that is perfectly flat. A step called twice with the same gradient doubles the update. Print the parameter norm before and after one step when a run refuses to learn.

The canonical loop

from torch.utils.data import DataLoader

loader = DataLoader(train_set, batch_size=64, shuffle=True, num_workers=4)

for epoch in range(epochs):
    model.train()
    for inputs, targets in loader:
        optimizer.zero_grad()
        logits = model(inputs)
        loss = criterion(logits, targets)
        loss.backward()
        optimizer.step()

    model.eval()
    with torch.no_grad():
        validate(model, val_loader)

Evaluation is a different mode

Dropout and batch normalisation behave differently in training and evaluation. model.eval() switches them. Forget it, and the validation number depends on which units happened to be dropped. torch.no_grad() switches off the graph and saves the activation memory that backward would have kept. Both are one line, and both are easy to omit.

Metrics need to be accumulated correctly. A mean of batch means is not the mean over the dataset when the last batch is smaller. Multiply each batch metric by its batch size, sum the results, and divide by the total number of examples. The same rule applies to accuracy, loss and anything else you average.

A scheduler adds one more call, and its placement matters. A scheduler is a rule that changes the learning rate during training. A per-epoch scheduler steps after the epoch. A per-batch scheduler steps after every optimiser update. Calling a per-epoch scheduler inside the batch loop decays the rate a thousand times too fast. The run then looks like it has a learning rate that is far too small.

What the loop hides

The DataLoader, the object that builds batches, can use worker processes. The order of the batches is then deterministic only if you seed the workers as well. It can pin memory, which locks pages so the GPU can read them faster, and it can move tensors asynchronously. A .item() call on a GPU tensor synchronises the device. Calling it every step can dominate the runtime of a small model, so log every N steps instead.

Checkpointing saves more than the weights. The optimiser state, the scheduler state and the epoch number are all needed to resume exactly. Saving only the model gives a run that restarts with fresh momentum and a learning rate at its initial value. That is fine for inference and wrong for a long training job.

The debugging order is fixed. Overfit one batch, then a small subset, then the full set. If one batch cannot reach a low loss, the bug is in the model or the loop. If it can, the data pipeline or the generalisation is the problem. This order saves more time than any hyperparameter search. A hyperparameter is a setting you choose before training, such as the learning rate.