D4.5 · LESSON 5 OF 6

The training loop

Translate an inspected NumPy model into tensors, modules, automatic differentiation and a clear training loop.

THE IDEA

The five lines every training loop contains

Inside the loop over batches, a PyTorch training step is the same five lines in the same order. Each one maps to something from the earlier chapters.

Line 1 of 5

▶ optimizer.zero_grad()
  pred = model(xb)
  loss = loss_fn(pred, yb)
  loss.backward()
  optimizer.step()

Clear the gradient left from the previous batch.

w.grad = None

One iteration on a one-weight model. Step through the lines and watch w and its gradient change.

real PyTorch (run it locally)

for xb, yb in loader:
    optimizer.zero_grad()          # 1. clear old gradients
    pred = model(xb)               # 2. forward
    loss = loss_fn(pred, yb)       # 3. loss
    loss.backward()                # 4. gradients
    optimizer.step()               # 5. update the weights

Two switches that are easy to confuse

model.train() and model.eval() change how certain layers behave (dropout is active only in train mode, batch norm uses batch or running statistics). torch.no_grad() stops gradient recording. They are separate settings: for validation you normally want both.

Predict first. Does model.eval() alone disable autograd?

Chapter context

PyTorch provides tensor operations and differentiation infrastructure. The framework removes bookkeeping, but it does not remove the need to understand dimensions, parameter ownership, device placement and evaluation behavior.

TRY IT YOURSELF

Part 1. Do it hands-on

PyTorch itself cannot run in a browser here, so the Python parts use a small PyTorch-style library that is already loaded. It has the same names and call shapes (torch.tensor, nn.Linear, loss.backward(), optimizer.step(), DataLoader) and the same behaviour for everything these lessons teach: shapes, accumulating .grad, graphs freed after backward(), parameter registration. It is not PyTorch: initial values and random numbers differ, and "cuda" is only a label. Every code block marked real PyTorch is for you to run on your own machine.

Four loops, each with exactly one problem. Read the symptom and the code, then name what is wrong.

Loop A

for xb, yb in loader:
    pred = model(xb)
    loss = loss_fn(pred, yb)
    loss.backward()
    optimizer.step()

Symptom: The loss improves for a few steps, then gets erratic or explodes.

Loop B

for xb, yb in loader:
    optimizer.zero_grad()
    pred = model(xb)
    loss = loss_fn(pred, yb)
    optimizer.step()

Symptom: The loss never changes at all, epoch after epoch.

Loop C

total = 0.0
for xb, yb in val_loader:
    pred = model(xb)
    total += loss_fn(pred, yb).item()

Symptom: The model contains Dropout. Validation loss is noisy and different every time you run it on the same data.

Loop D

pred = model(xb)        # shape (32, 1)
loss = loss_fn(pred, yb)  # yb has shape (32,)

Symptom: The loss is a surprising number, and a warning about target size appears.

To do: Loop A: which line is missing or wrong?
To do: Loop B: which line is missing or wrong?
To do: Loop C: which line is missing or wrong?
To do: Loop D: which line is missing or wrong?

Part 2. Do it in code

train(model, X, y, lr, epochs, batch_size): X and y are tensors. Build loader = DataLoader(TensorDataset(X, y), batch_size=batch_size, shuffle=False), loss_fn = nn.MSELoss() and optimizer = torch.optim.SGD(model.parameters(), lr=lr). For every epoch loop over the loader doing the standard five steps, and record the mean batch loss of the epoch. Return the list of epoch losses.
evaluate(model, X, y): switch the model to eval mode, compute the MSE on all of X inside torch.no_grad(), restore the mode the model was in before, and return the loss as a Python float.

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

def train(model, X, y, lr, epochs, batch_size):
    pass

def evaluate(model, X, y):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Does model.eval() alone disable autograd?

ANoBYes

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

For a batch X of shape (B,d), a dense output with h features has shape (B,h). torch.nn.Linear stores weight with shape (h,d) and computes X times weight-transpose plus bias. Parameter count is d*h+h. PyTorch execution belongs in a linked notebook or supported runtime; the chapter does not assume browser PyTorch.

Worked example

A batch has shape (2,4). A Linear(4,3) layer stores a (3,4) weight and a length-three bias. Its output is (2,3). The parameter count is 4×3+3=15. This differs from the batch's element count: model parameters are shared across both examples.

Reference code

The following example uses the Python standard library and does not call a remote service.

batch, inputs, outputs = 2,4,3
weight_shape = (outputs,inputs)
output_shape = (batch,outputs)
assert weight_shape == (3,4)
assert output_shape == (2,3)
assert inputs*outputs+outputs == 15

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

eval mode and no_grad are independent. Model parameters do not multiply with batch size. A tensor detached for logging must not replace the tensor used to compute the training loss.