Training Loops in PyTorch

This article studies a single function, train_classifier_simple, which is used to fine-tune a model to classify spam. Its job is to train the model on a training dataset and, as training proceeds, to evaluate it on a validation dataset. Reading the function as a program is the best way in: the function is driven by a few inputs, and it hands a set of records back to its caller. Once you can see the shape of the loops, the details of each step fall into place.

Before tracing the loops, look at what the function receives and what it returns. The function takes the model, the two data loaders, an optimizer, the device to run on, and four control values: num_epochs, eval_freq, and eval_iter. The control values set how long training runs and how often the model is checked. The function returns four lists and a counter.

The returned lists let you follow training progress over time. Because examples_seen is returned alongside them, each recorded loss or accuracy can be matched to how much training data the model had consumed at that point.

The function uses two nested loops. The outer loop, written as for epoch in range(num_epochs), runs once per epoch. At the start of each epoch the function calls model.train(), which puts the model in training mode. The inner loop, written as for input_batch, target_batch in train_loader, runs once per batch of training data. Each batch passes through the same sequence of steps. Only after the inner loop has finished every batch does the outer loop advance to the next epoch.

The diagram below shows the order of these loops. Trace it with a concrete picture in mind: if the training loader contains a fixed number of batches, the inner body runs that many times in every epoch, and the whole run executes the inner body once for each batch in each of the num_epochs passes.

Each batch runs four steps in a fixed order. The order matters because each step consumes what the previous one produced. The steps are listed below, followed by a skeleton that shows where each one sits inside the batch body.

The flowchart below follows a single batch from the top of the batch body to the point where the bookkeeping begins. Notice that the loss is the output of one step and the input to the next, so the sequence is a chain rather than a set of independent actions.

The reset sits at the top of the batch body for a specific reason. The source describes optimizer.zero_grad() as clearing the gradients left over from the previous batch iteration. Gradients from one batch must not be mixed into the calculation for the next batch, so the reset happens before loss.backward() produces new gradients for the current batch.

After the weight update, the function does two bookkeeping actions. It adds the number of rows in the current input batch to examples_seen, using input_batch.shape[0]. It then increments global_step by one. Because examples_seen counts examples rather than tokens, the source notes this as a change from the earlier pretraining code.

The two counters measure different things. examples_seen tracks how much data the model has consumed, while global_step tracks how many batches have been processed. Only global_step is used to decide whether an evaluation happens.

The check if global_step % eval_freq == 0 controls how often the model is evaluated. The % operator gives the remainder after division, so the condition is true whenever global_step is an exact multiple of eval_freq. This makes eval_freq a count of training steps between evaluations, not a count of epochs.

The counter global_step starts at -1 and increases with every batch across all epochs. It is never reset at the start of an epoch, so the count keeps running through the whole training run. An evaluation therefore can fire in the middle of an epoch, and the next batches continue normally afterward.

With eval_freq set to 5, the check passes at global_step values that are multiples of 5, such as 5, 10, and 15, and the evaluation runs at those points. Between them, the batches only train. Changing eval_freq changes the spacing of these points without changing the training itself.

When the check passes, the function calls evaluate_model to obtain a training loss and a validation loss, appends both to their lists, and prints them with the current epoch and step. It then calculates accuracy on both loaders and records those values. The accuracy calculation uses num_batches=eval_iter, so only a limited number of batches from each loader are used for this measurement.

The evaluate_model function is the other half of the picture. The source states that it is identical to the one used for pretraining. It switches the model into evaluation mode with model.eval(), and it runs its calculations inside a torch.no_grad() block so that no gradients are computed. It estimates loss on both datasets and does not update any weights. The comparison below sets it against a training step.

The function ends by setting the model back to training mode with model.train(). This matters because train_classifier_simple calls evaluate_model partway through an epoch. Without that final call, the next batches would run with the model still in evaluation mode. This is also the reason the training function calls model.train() at the start of each epoch: the mode is set again explicitly rather than assumed.