Training Neural Networks

An epoch runs all training batches once, but the meaningful unit of optimisation is usually the update step rather than the epoch. In training mode, each batch follows the complete loop—forward pass, loss, backward, and optimiser step—while validation measures held-out data without updating parameters. model.train() enables training behaviour such as dropout and current batch statistics; model.eval() switches layers to inference behaviour, and torch.no_grad() prevents autograd from retaining unnecessary activations. Save the checkpoint with the best validation metric rather than automatically trusting the final epoch.

Validation loss is the signal for early stopping, model selection, and learning-rate scheduling. Training loss alone cannot tell you whether the model is learning a reusable pattern or memorising. A widening train/validation gap is evidence to investigate capacity, augmentation, data quality, or regularisation—not a cue to simply train longer.

The canonical PyTorch loop is six lines inside two nested loops: for each epoch, for each batch, move the batch to the device, zero the gradients, compute the loss from a forward pass, call backward(), then step() the optimiser (and the scheduler, if it is per-step). The order matters in exactly one place: zero_grad must come before backward, and backward before step. Everything else in a real training script — logging, checkpointing, validation, clipping, mixed precision — hangs off this skeleton without changing it.

PyTorch makes the loop explicit rather than hiding it behind a fit() call, which is the reason it is the research default: every intermediate value is inspectable, so any of the failures in this module can be diagnosed where it happens. The cost is that the six lines are yours to get right, and two orderings — a missing zero_grad, a step before backward — produce runs that look plausible and learn wrongly or not at all.

A step (or iteration, or update) is one optimiser update: one batch forward, backward, and step(). An epoch is one complete pass over the training set, so it contains ceil(len(dataset) / batch_size) steps — which is what len(dataloader) returns. Total steps for a run is epochs × steps_per_epoch. With gradient accumulation over N micro-batches, one optimiser step spans N forward/backward passes, so steps and batches stop being the same thing.

Almost every training hyperparameter is defined against one of these units, and mixing them up breaks things silently. A cosine schedule with T_max in steps, stepped once per epoch, finishes its decay in the first few percent of the run. A warmup of '500' means 500 steps, not 500 epochs. And comparing two runs by epoch count is meaningless if their batch sizes differ, because the model that took eight times as many updates has done eight times as much learning per epoch.

Plot both losses against epochs on the same axes. A healthy run has both falling, with the validation curve slightly above the training one and flattening out. An overfitting run has training loss continuing to fall while validation loss bottoms out and then turns upward — the model is still improving its fit to the training set, and that improvement has stopped transferring. The epoch where validation loss turns is the checkpoint worth keeping.

This one plot answers 'should I train longer, add capacity, or add data?' before you spend anything. Both curves high and flat means underfitting: more capacity, a higher learning rate, or longer training. Training loss near zero with validation loss rising means overfitting: more data, augmentation, regularisation, or an earlier stop. A validation curve that is wildly noisier than the training curve usually means the validation set is too small to measure with, not that the model is unstable. The ML syllabus's bias-variance module makes the same argument across model capacities; here you watch it happen within one run.

A resumable checkpoint is a dict: the model's state_dict, the optimizer's state_dict, the scheduler's state, the epoch or step count, and the best metric so far. Loading it means constructing the same model and optimizer objects and calling load_state_dict on each. Saving only model.state_dict() is enough to *serve* a model and not enough to *resume* one.

The optimizer holds real state. Adam's two moment estimates per parameter took hundreds of steps to build up, and resuming without them restarts from zero moments — which produces a visible loss spike and, with bias correction restarting too, several hundred steps of degraded updates. The scheduler's step count matters just as much: without it, a resumed run restarts its warmup or its cosine decay, so the learning rate is wrong for the rest of training. Both failures look like 'the resume made the model worse', which is exactly what happened.

Early stopping watches the validation metric and halts when it has not improved for patience consecutive evaluations, keeping the best checkpoint rather than the last. Three parameters define it: the metric (validation loss, or the metric you actually care about), the patience (how many non-improving evaluations to tolerate), and the minimum improvement that counts as one. It is the deep-learning version of choosing model complexity, with epochs as the complexity knob.

Deep networks will happily continue to a training loss of zero, so 'when to stop' is a real decision rather than a technicality, and the loss curve above shows that stopping in the wrong place costs measurable accuracy. Patience matters because validation curves are noisy: patience of 1 stops on the first unlucky evaluation, while a patience of 5 to 10 evaluations rides out the noise and still catches a genuine turn. Restoring the best weights is the part people forget — stopping late and keeping the final weights recovers none of the benefit.

Log at least five things per step or per epoch: the training loss, the validation metric, the current learning rate (read from opt.param_groups[0]["lr"], not from your config), the global gradient norm before clipping, and a couple of actual sample predictions alongside their targets. Wall-clock time per step and peak memory are worth adding, because they are what tells you whether the data loader or the model is the bottleneck.

The loss is one number and most failures have a distinctive signature in the others. A schedule that never advanced shows up in the logged learning rate and nowhere else. Vanishing or exploding gradients show up in the gradient norm several hundred steps before the loss does anything visible. And a model achieving a plausible loss while predicting nonsense — the same class for everything, or a shifted-by-one sequence — is only ever caught by looking at predictions, which is why a handful of real examples in the log is worth more than another scalar.

A NaN loss means a non-finite number entered the computation and then spread: any arithmetic involving NaN produces NaN, so one bad value contaminates every parameter on the next backward pass and the run never recovers. The usual causes, in rough order of frequency: a learning rate high enough that one step overflows the weights; a log(0) or divide-by-zero from computing a loss by hand instead of using the logit-based version; NaNs or infinities already present in the input data; and mixed-precision overflow in float16 without a working gradient scaler.

It is unrecoverable and it is cheap to localise, so the value is in knowing the checklist rather than in bisecting for an hour. The decisive question is whether the input was already bad: check the batch for non-finite values first, because if the data is clean the problem is numerical and the learning rate is the first thing to halve. Adding a per-step finite check costs nothing and turns a silent run of NaNs into an exception at the exact step and batch that caused it.