Chapter 9: Models and the Training Loop

Chapter 8 worked out one gradient and stepped one weight. A real model has thousands of weights and learns from a batch of examples at a time, and the code that does it is the training loop. It is five lines long. This chapter is about reading them, and about the bugs that run without an error.

A model is an nn.Module: it makes its layers in __init__ and uses them in forward. nn.Linear multiplies by weights and adds a bias, and nn.ReLU zeroes the negatives between two of them. nn.Sequential chains layers when that’s all a model does. nn.Embedding turns integer ids, such as a language model’s tokens, into rows of learned numbers.

import torch
from torch import nn
# A model is a class: layers made in __init__, used in forward.
class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(2, 4)
        self.out = nn.Linear(4, 2)

    def forward(self, x):
        return self.out(torch.relu(self.hidden(x)))
model = Classifier()
# Call the model, not forward: model(x) runs forward for you.
model(torch.randn(5, 2)).shape

torch.Size([5, 2])

# The same model as a chain of layers.
nn.Sequential(nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 2))

Sequential(
  (0): Linear(in_features=2, out_features=4, bias=True)
  (1): ReLU()
  (2): Linear(in_features=4, out_features=2, bias=True)
)

# Every weight and bias the optimizer will step: 2*4 + 4 + 4*2 + 2.
sum(p.numel() for p in model.parameters())

22

# An embedding looks rows up by id: 10 ids, 3 numbers each.
nn.Embedding(10, 3)(torch.tensor([7, 2, 7])).shape

torch.Size([3, 3])

A loss function scores the output. nn.CrossEntropyLoss takes raw scores and class indices; softmax the scores first and you get a wrong loss, not an error. nn.MSELoss is for predicting numbers. The optimizer holds the model’s weights: torch.optim.SGD steps each one by the learning rate times its gradient, and torch.optim.AdamW sizes the step for each weight.

import torch
from torch import nn
# Raw scores (logits) for 3 examples and 2 classes, and each one's class.
logits = torch.tensor([[2.0, -1.0], [0.5, 0.5], [-1.0, 3.0]])
target = torch.tensor([0, 1, 1])
nn.CrossEntropyLoss()(logits, target)

tensor(0.2533)

# Classes must be integers; float labels raise.
nn.CrossEntropyLoss()(logits, target.float())

RuntimeError: expected target dtype to be Long or Byte, but got Float

# Softmaxing first doesn't raise; it just gives a different, wrong loss.
nn.CrossEntropyLoss()(logits.softmax(dim=1), target)

tensor(0.4520)

# For numbers rather than classes: the mean squared error.
nn.MSELoss()(torch.tensor([2.5, 0.0]), torch.tensor([3.0, 0.0]))

tensor(0.1250)

A Dataset hands out one example by index, and a DataLoader groups them into batches. It keeps their order unless you pass shuffle=True, which you want for training. One pass over all of them is an epoch.

import torch
from torch.utils.data import DataLoader, TensorDataset
X = torch.arange(10.0).reshape(5, 2)
y = torch.tensor([0, 1, 0, 1, 1])
dataset = TensorDataset(X, y)
# A Dataset hands out one example by index.
len(dataset), dataset[0]

(5, (tensor([0., 1.]), tensor(0)))

# A DataLoader groups examples into batches, in order unless you pass shuffle=True.
for xb, yb in DataLoader(dataset, batch_size=2):
    print(xb.shape, yb.tolist())

torch.Size([2, 2]) [0, 1]
torch.Size([2, 2]) [0, 1]
torch.Size([1, 2]) [1]

The stepper runs one epoch of a tiny classifier: 16 points, two batches of 8. Each batch goes round the same five steps: zero_grad, forward, the loss, backward and optimizer.step(). In this run, after each step the same batch scores a lower loss.

import torchfrom torch import nnfrom torch.utils.data import DataLoader, TensorDatasettorch.manual_seed(0)X = torch.randn(16, 2)y = (X.sum(dim=1) > 0).long()loader = DataLoader(TensorDataset(X, y), batch_size=8, shuffle=True)model = nn.Sequential(nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 2))loss_fn = nn.CrossEntropyLoss()optimizer = torch.optim.SGD(model.parameters(), lr=0.5)model.train()for xb, yb in loader:    optimizer.zero_grad()    logits = model(xb)    loss = loss_fn(logits, yb)    loss.backward()    optimizer.step()

2

This is where most people slip. Leave out zero_grad and nothing raises: each backward adds to the gradients of every batch before it, as in Chapter 8. Move it between backward and optimizer.step() and nothing raises either: the step skips every weight whose .grad is None, so nothing learns. Keep it first in the loop.

model.train() and model.eval() switch layers that act differently in training, such as nn.Dropout, which zeroes random values. Forget model.eval() before scoring a model with dropout and the score is noisy and can change from call to call. It doesn’t stop autograd recording, so score inside torch.no_grad too.

import torch
from torch import nn
model = nn.Sequential(nn.Linear(4, 4), nn.Dropout(p=0.5))
x = torch.ones(1, 4)
# A fresh model is in training mode: dropout zeroes each value with probability 0.5, afresh each call.
model.training

True

torch.equal(model(x), model(x))

False

# Evaluation mode turns dropout off, so the same input gives the same output.
model.eval()

Sequential(
  (0): Linear(in_features=4, out_features=4, bias=True)
  (1): Dropout(p=0.5, inplace=False)
)

torch.equal(model(x), model(x))

True

# eval() doesn't stop autograd recording; torch.no_grad() does.
model(x).requires_grad

True

with torch.no_grad():
    print(model(x).requires_grad)

False

A model’s state_dict is its weights by name. torch.save writes it to a file and torch.load reads it back. Since PyTorch 2.6, torch.load loads only tensors and plain Python values by default, so save the state_dict, not the whole model.

import tempfile
from pathlib import Path
import torch
from torch import nn
model = nn.Sequential(nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 2))
# The state_dict: every weight by name. Numbers 0 and 2 are the two Linear layers.
{name: tuple(t.shape) for name, t in model.state_dict().items()}

{'0.weight': (4, 2), '0.bias': (4,), '2.weight': (2, 4), '2.bias': (2,)}

# path is a model.pt file in a scratch folder.
torch.save(model.state_dict(), path)
# Load into a fresh model of the same shape.
fresh = nn.Sequential(nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 2))
fresh.load_state_dict(torch.load(path))

<All keys matched successfully>

x = torch.randn(3, 2)
torch.equal(fresh(x), model(x))

True

The last bug does raise: a model and a batch on different devices. An nn.Module’s .to(device) moves its weights in place, but every batch needs its own .to(). This page runs on the CPU, so a dtype mismatch stands in, fixed the same way. On a GPU, mixed precision runs some of the math in 16-bit floats to save time and memory.

import torch
from torch import nn
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = nn.Linear(2, 1)
# A module's .to() moves its weights in place and returns the module itself.
model.to(device) is model

True

# Every batch has to follow it; a mismatch raises. This page runs on the CPU,
# so a float64 batch stands in for one left on the wrong device.
model(torch.ones(3, 2, dtype=torch.float64))

RuntimeError: mat1 and mat2 must have the same dtype, but got Double and Float

model(torch.ones(3, 2, dtype=torch.float64).to(device, torch.float32)).shape

torch.Size([3, 1])

The habit this chapter adds: read a training loop for its five lines in order, zero_grad first, and switch modes on purpose, model.train() to train and model.eval() to score. How Do You Train One Model on 16,000 GPUs? runs this same loop on thousands of GPUs at once, splitting the batch and the model between them.

Homework

The packet has two parts: fix three bugs in a broken training loop (a missing zero_grad, a score taken without model.eval(), and a dtype mismatch standing in for a device one), then build and train a small classifier to at least 90% accuracy on points it never saw, and load its state_dict back. Every exercise runs on the CPU; the README says how to borrow a free GPU in Google Colab if you want to try one.