The training loop

Every step runs the same five parts

1. Get a batch. A data loader hands over the next minibatch, already prepared by worker processes, often with random changes such as rotations (data augmentation). The loop moves it to the GPU (tensors on the GPU).
2. Forward pass and loss. The model turns the batch into predictions, and the loss scores them with one number. The model's arithmetic often runs in 16-bit numbers to save time and memory (mixed precision). The loss itself is usually kept in 32 bits.
3. Backward pass. Backpropagation computes the gradient of the loss for every weight; in PyTorch, loss.backward() does it (automatic gradients). With several GPUs, the GPUs average their gradients here.
4. Take a step. Gradient clipping shrinks the gradients if they are unusually large. The optimizer, often AdamW, updates the weights, a refinement of plain gradient descent. The learning-rate schedule sets the step size for the next step.
5. Update the averaged weights. A slowly moving copy of the weights takes a small step toward the new weights (a moving average of the weights).

And now and then, for example every 1,000 steps:

6. Check and save. Measure the loss on the validation set, which the model never trains on, and write a checkpoint so that training can resume after a stop.

Steps, epochs and checks run on different clocks

One pass of parts 1 to 5 is a step. One pass through all the training examples is an epoch. With 1,000 examples and batches of 50, an epoch is 20 steps, and 10 epochs are 200 steps. Checks and saves are usually counted in steps: every 50 steps in the code below, every few thousand in a long run.

All six parts fit in about 30 lines

This PyTorch script trains a small network on made-up data on the CPU. The numbered comments match the parts above. Each line is explained in the unit linked from its part.

loop.py

import copy
import torch
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(0)                                    # same random numbers on every run
x = torch.randn(1200, 8); y = x.sum(1, keepdim=True) + 0.1 * torch.randn(1200, 1)
train_dl = DataLoader(TensorDataset(x[:1000], y[:1000]), batch_size=50, shuffle=True)
val_x, val_y = x[1000:], y[1000:]                       # held out: never trained on

model = torch.nn.Sequential(torch.nn.Linear(8, 32), torch.nn.ReLU(), torch.nn.Linear(32, 1))
ema = copy.deepcopy(model)
opt = torch.optim.AdamW(model.parameters(), lr=0.01)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=200)

step = 0
for epoch in range(10):                                 # 20 batches per epoch
    for xb, yb in train_dl:                             # 1. get a batch
        with torch.autocast("cpu", dtype=torch.bfloat16):
            loss = ((model(xb) - yb) ** 2).mean()       # 2. forward pass and loss
        opt.zero_grad()
        loss.backward()                                 # 3. gradients
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step(); sched.step(); step += 1             # 4. step, schedule
        with torch.no_grad():                           # 5. averaged copy
            for pe, p in zip(ema.parameters(), model.parameters()):
                pe.lerp_(p, 0.05)                       # beta = 0.95
        if step % 50 == 0:                              # 6. now and then
            with torch.no_grad():
                val = ((ema(val_x) - val_y) ** 2).mean().item()
            print(f"step {step}: train loss {loss.item():.3f}, validation loss {val:.3f}")
            torch.save({"model": model.state_dict(), "ema": ema.state_dict(), "step": step,
                        "optimizer": opt.state_dict(), "schedule": sched.state_dict()}, "ckpt.pt")

step 50: train loss 0.067, validation loss 0.594
step 100: train loss 0.027, validation loss 0.048
step 150: train loss 0.027, validation loss 0.030
step 200: train loss 0.019, validation loss 0.026

TensorDataset pairs each row of x with the same row of y. CosineAnnealingLR lowers the learning rate by cosine decay, to 0 at step T_max = 200. A semicolon separates two statements on one line. The % sign gives the remainder of a division, so step % 50 == 0 is true at every 50th step.

The validation loss uses the averaged copy, because that is the model the run keeps. At step 50 it is still high, because the copy averages over about the last 20 steps and lags behind the fast early progress. (The trained weights alone score about 0.056 there.) By step 150 the copy has caught up. pe.lerp_(p, 0.05) is the averaging rule with β = 0.95, written as "move 5% of the way toward p".