12.5.3 — The Training Loop in PyTorch

12.5.2 did one forward pass, one backward pass and one update, on one example. Training Cartwise's photo checker means doing that hundreds of thousands of times, on batches of photos, while watching a validation set to know when to stop. The code that does this is called the training loop, and almost every deep learning project in the world, from a two-layer network to a large language model, has the same five lines at its centre.

This page builds that loop in PyTorch line by line, then opens up the one part of it that hides the most important decisions: the optimiser, which decides exactly how each weight moves.

1. The four pieces, before the loop

A PyTorch training script always assembles the same four things before any training happens.

The data, served in batches. PyTorch separates what the data is from how it is served. A Dataset knows how many examples there are and how to fetch example number i. A DataLoader wraps it and hands out shuffled mini-batches, from 12.2.4, one at a time.

The model. A class that inherits from nn.Module, which creates its layers in __init__ and says how data flows through them in forward.

The loss. A function that compares the model's output with the right answers and returns one number.

The optimiser. An object that holds references to all the model's weights and knows how to update them from their gradients.

Here is all four for the photo checker. To keep the example runnable anywhere, each photo has already been turned into 512 numbers by a pre-trained image network (the transfer learning of 12.1.2; 12.5.5 shows how), and our network learns the "bad photo or not" decision on top.

import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader

torch.manual_seed(0)                                       # (1) repeatable randomness

# Stand-in data: 4,000 photos × 512 features, label 1 = bad photo
X = torch.randn(4_000, 512)
y = (X[:, :8].sum(dim=1) > 1.5).float()                    # (2) a hidden rule to learn
train_ds = TensorDataset(X[:3_200], y[:3_200])
val_ds   = TensorDataset(X[3_200:], y[3_200:])
train_dl = DataLoader(train_ds, batch_size=64, shuffle=True)   # (3)
val_dl   = DataLoader(val_ds, batch_size=256)

class PhotoChecker(nn.Module):                             # (4)
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(512, 128), nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 1),                             # one raw score (a logit)
        )
    def forward(self, x):
        return self.net(x).squeeze(-1)                     # shape (batch,)

model = PhotoChecker()
loss_fn = nn.BCEWithLogitsLoss()                           # (5)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)   # (6)

Line (1) fixes the random number generator, so the starting weights, the stand-in data and the shuffling are the same every run. Without it, two runs of the same code give different results, which makes debugging miserable. Line (2) invents labels from a hidden rule (a photo is "bad" when its first eight features add up past 1.5), so that there is something real to learn; in the real project, X and y would come from the labelled photos.

Line (3) creates the loaders. batch_size=64 means each step sees 64 photos. shuffle=True reorders the training photos every epoch, so batches are different each time; the validation loader does not shuffle, because order does not matter when you only measure.

Line (4) defines the model: 512 inputs, one hidden layer of 128 ReLU neurons, a dropout layer (explained in 12.5.4), and one output. Notice that there is no sigmoid at the end. The model outputs a raw logit, and that is deliberate, as the next line explains.

Line (5) is the loss. BCEWithLogitsLoss takes raw logits and applies the sigmoid and the binary cross-entropy together, in one carefully written function. Doing the two steps separately can produce log(0) when a logit is very large, because the sigmoid rounds to exactly 1.0; the combined version uses the log-probability trick from 12.2.6 and never does. The multi-class equivalent, nn.CrossEntropyLoss, likewise takes raw logits and applies softmax itself. Putting your own softmax or sigmoid in front of these is one of the most common bugs in deep learning code: it applies the squashing twice and quietly cripples training.

Line (6) is the optimiser. model.parameters() hands it every weight and bias in the model. AdamW and its two settings are the subject of section 3.

2. The loop

def evaluate(model, loader):
    model.eval()                                           # (a) switch off dropout
    total, correct, loss_sum = 0, 0, 0.0
    with torch.no_grad():                                  # (b) no gradient bookkeeping
        for xb, yb in loader:
            logits = model(xb)
            loss_sum += loss_fn(logits, yb).item() * len(yb)
            correct += ((logits > 0) == (yb == 1)).sum().item()
            total += len(yb)
    return loss_sum / total, correct / total

for epoch in range(10):
    model.train()                                          # (c) switch dropout back on
    for xb, yb in train_dl:
        opt.zero_grad()                                    # (1) clear old gradients
        logits = model(xb)                                 # (2) forward
        loss = loss_fn(logits, yb)                         # (3) measure
        loss.backward()                                    # (4) backpropagate
        opt.step()                                         # (5) update every weight
    val_loss, val_acc = evaluate(model, val_dl)
    print(f"epoch {epoch}  val_loss {val_loss:.3f}  val_acc {val_acc:.3f}")

The five numbered lines are the heart of every training loop ever written.

Line (1), zero_grad(), clears the gradients left over from the previous step. This surprises everyone the first time: PyTorch adds new gradients to whatever is already stored in .grad; it does not replace them. That design is deliberate, because adding up gradients over several batches before updating is a useful trick (gradient accumulation, in 12.5.7). But it means that if you forget this line, each step's update uses the sum of every gradient since training began, and training goes haywire.

Line (2) runs the batch of 64 photos through the model, producing 64 logits. Line (3) compares them with the 64 labels and returns the average loss over the batch. Line (4) runs backpropagation exactly as in 12.5.2, filling every weight's .grad. Line (5) asks the optimiser to move every weight using its gradient.

The lettered lines handle the difference between training and measuring. Line (c), model.train(), and line (a), model.eval(), switch layers that behave differently in the two situations. Dropout randomly switches off neurons during training and must stop doing so when you measure; batch normalisation, in 12.5.4, also behaves differently. Forgetting model.eval() makes validation scores noisy and slightly wrong. Line (b) turns off gradient recording during evaluation, saving memory and time, as 12.5.2 explained.

Run it and the validation loss falls over the epochs while accuracy rises, typically ending at well over 90 percent on this stand-in data (the exact numbers depend on the library version and machine). Two words recur in every training log, so to be precise: each pass through the five lines is one step; one pass through all 3,200 training photos, here 3{,}200 / 64 = 50 steps, is one epoch.

3. The optimiser: how each weight actually moves

Line (5), opt.step(), is where the most consequential choices hide. Plain gradient descent from 12.2.4 moves each weight by the learning rate times its gradient. That works, but two problems make it slow on real networks, and the optimisers below were invented to fix them.

Problem one: zigzagging in narrow valleys, and momentum

12.2.3 showed that a loss landscape is often a long, narrow valley: steep across, gentle along. Plain gradient descent points mostly across the valley, the steep direction, so it bounces from wall to wall while creeping slowly along the floor.

A narrow valley seen from above. Plain gradient descent (red) zigzags because each step points mostly across the valley. Momentum (green) keeps a running average of past steps: the back-and-forth components cancel out, and the steady forward components build up.

Momentum keeps a running total of recent gradients and steps along that total instead of the latest gradient alone:

\mathbf{v} \leftarrow \beta\,\mathbf{v} + \nabla L, \qquad \mathbf{w} \leftarrow \mathbf{w} - \eta\,\mathbf{v}

Read it as "the velocity becomes beta times the old velocity plus the new gradient; the weights move by the learning rate times the velocity". The number \beta (beta), typically 0.9, sets how long the memory is. The image is a heavy ball rolling downhill: it does not stop and turn at every bump, it carries its speed. In a narrow valley, the across-the-valley parts of successive gradients point alternately left and right and cancel out in the running total, while the along-the-valley parts all point the same way and add up. The ball speeds up along the floor.

It is worth seeing that momentum is not free. Here is the simple bowl L(w) = (w - 3)^2 from 12.2.4, starting at w = 5 with learning rate 0.1, with and without \beta = 0.9:

Step	Plain: w	Momentum: gradient	Momentum: velocity	Momentum: w
1	4.600	4.000	4.000	4.600
2	4.280	3.200	6.800	3.920
3	4.024	1.840	7.960	3.124
4	3.819	0.248	7.412	2.383
5	3.655	−1.234	5.436	1.839

By step 3, momentum has almost reached the answer, far ahead of plain descent at 4.024. But the ball has built up speed, and at step 4 it rolls straight past the bottom at 3 and keeps going, because its velocity is still large even though the gradient has nearly vanished. It will swing back and settle eventually. On a round bowl like this one, momentum overshoots; in the long narrow valleys of real networks, its smoothing is exactly what you want, and SGD with momentum 0.9 is still the standard for training many image networks.

Problem two: one learning rate for every weight, and Adam

Different weights see gradients of wildly different sizes. Weights connected to a rarely active feature get small, infrequent gradients; weights in a busy layer get large, constant ones. One shared learning rate is too big for some and too small for others.

Adam, published by Diederik Kingma and Jimmy Ba in 2014 (the name comes from adaptive moment estimation), gives every weight its own effective step size. For each weight it keeps two running averages: of the gradient itself, m, which is momentum, and of the squared gradient, s, which measures how large that weight's gradients typically are. It then divides one by the square root of the other:

m \leftarrow \beta_1 m + (1 - \beta_1)\, g, \qquad s \leftarrow \beta_2 s + (1 - \beta_2)\, g^2, \qquad w \leftarrow w - \eta \, \frac{\hat{m}}{\sqrt{\hat{s}} + \epsilon}

Read the last part as "the weight moves by the learning rate times the average gradient divided by the typical gradient size". The defaults are \beta_1 = 0.9, \beta_2 = 0.999, and a tiny \epsilon (epsilon, about 10^{-8}) to avoid dividing by zero. The hats on \hat{m} and \hat{s} mark a small bias correction: both averages start at zero, so for the first few steps they are too small, and dividing by 1 - \beta^t (where t is the step number) scales them up to the right size.

The division is the key idea. Run Adam on our bowl from w = 5 with \eta = 0.1: the first step moves the weight by exactly 0.1, to 4.9, then 4.800, 4.701, 4.602. On step 1, \hat{m} = g and \sqrt{\hat{s}} = |g|, so the ratio is exactly 1 in size, whatever the gradient is. Now make every gradient a thousand times larger, as if this weight lived in a much busier layer. Plain gradient descent would take steps a thousand times bigger and fly off. Adam's path is identical: 4.9, then 4.800, because the thousand appears on the top and the bottom of the ratio and cancels. Adam's step size is set by the learning rate, not by the size of the gradient. Every weight moves at a sensible pace whether its gradients are tiny or huge, which is why Adam works well with little tuning and is the default for most deep learning, including all transformers.

Weight decay, and why AdamW exists

12.3.3 introduced L2 regularisation: add a penalty for large weights to the loss. In neural networks this is called weight decay, because its effect is to shrink every weight a little at every step. With plain gradient descent the two are the same thing. With Adam they are not: if the penalty is added to the loss, its gradient goes through Adam's division and gets scaled differently for every weight, which weakens the regularisation exactly where it is most needed. Ilya Loshchilov and Frank Hutter showed this in 2017 and proposed AdamW, which applies the shrinking directly to the weights, separately from the gradient step. AdamW with a weight decay around 0.01 to 0.1 is the standard choice for training transformers. Line (6) of the setup used it.

4. Learning rate schedules: changing the step size over time

A single learning rate is rarely best for the whole of training. Early on, the weights are random and the gradients are large and unreliable, so big steps can wreck the network. Late in training, the model is near a good solution and needs small steps to settle into it. A learning rate schedule changes \eta as training goes on. The most common shape today has two parts.

Warmup followed by cosine decay. The learning rate climbs from near zero to its peak over the first few hundred or thousand steps, then glides down along half a cosine wave to near zero by the end.

Warmup. For the first few hundred or few thousand steps, the learning rate climbs in a straight line from near zero to its chosen peak. This protects the network while its weights are random and Adam's running averages have not yet settled. Large models, transformers especially, often fail outright without it, with the loss shooting to nan in the first minutes.

Cosine decay. After warmup, the learning rate follows half a cosine wave down to near zero:

\eta_t = \eta_{\text{peak}} \cdot \tfrac{1}{2}\left(1 + \cos\!\left(\pi \cdot \frac{t - t_{\text{warm}}}{T - t_{\text{warm}}}\right)\right)

Read it as "the learning rate at step t is the peak times one half of one plus the cosine of pi times the fraction of the post-warmup steps completed". At the start of decay the fraction is 0, \cos 0 = 1, and the rate is the peak. At the end, the fraction is 1, \cos \pi = -1, and the rate is 0. In between it falls slowly at first, fastest in the middle, and gently at the end. Ilya Loshchilov and Frank Hutter, the same pair behind AdamW, popularised it in 2016.

PyTorch provides schedulers that adjust the optimiser's learning rate for you. You call scheduler.step() once per training step, after opt.step().

5. Choosing the numbers

Practical defaults that work for most projects, and the reasons for them:

• Optimiser: AdamW. Reach for SGD with momentum 0.9 if you are training image networks from scratch and following a published recipe that uses it.
• Learning rate: around 10^{-3} for small networks trained from scratch with Adam; around 10^{-5} to 10^{-4} when fine-tuning a large pre-trained model, because big steps would destroy what it already knows (12.5.7 explains this "catastrophic forgetting"). If unsure, try a few values spaced by factors of three and keep the largest that trains smoothly.
• Batch size: as large as fits in memory, within reason; 32 to 512 for most models. If you double the batch size, the gradient estimate is less noisy, and you can usually raise the learning rate somewhat too.
• Schedule: warmup of a few percent of the steps, then cosine decay.
• Epochs: enough that validation loss stops improving, with early stopping (12.5.4) to pick the best point.

Above all, change one thing at a time and log every run, which is the experiment tracking of 12.5.7. A training loop is an experiment, and an experiment you cannot reproduce teaches you nothing.

Recall

• Four pieces: a Dataset and shuffled DataLoader, an nn.Module model, a loss, an optimiser. Use BCEWithLogitsLoss or CrossEntropyLoss on raw logits; never add your own sigmoid or softmax in front.
• The training loop's five lines: zero_grad, forward, loss, backward, step. PyTorch adds gradients, so forgetting zero_grad breaks training. Use model.train() and model.eval(), and torch.no_grad() when measuring.
• Momentum averages recent gradients: zigzags cancel and steady directions build speed, at the cost of overshooting on simple bowls.
• Adam divides each weight's averaged gradient by its typical gradient size, so step size is set by the learning rate, not the gradient's scale. AdamW applies weight decay separately and is the standard for transformers.
• A learning rate schedule: linear warmup protects early training; cosine decay settles the model. Defaults: AdamW, about 10^{-3} from scratch and 10^{-5} to 10^{-4} for fine-tuning.

Self-test: What happens to training if opt.zero_grad() is left out, and why? On Adam's first step, why does the weight move by exactly the learning rate whatever the gradient? Why should a model trained with BCEWithLogitsLoss not end in nn.Sigmoid()?

Next: 12.5.4 Making Training Work — the loop is simple; getting a deep network to actually train well is not, and the next page covers the starting weights, normalisation, shortcuts, regularisers and debugging habits that make the difference.