The training loop

What you will be able to do

Train the ticket router for real: repeat five steps, in a fixed order, until the loss is low, and keep a log of it. Every training run in PyTorch is this loop, and a missing step usually gives no error message, only a network that learns badly or not at all.

NotebookDo this lesson in a notebookColab · Kaggle · GitHub

The notebook has every code cell, task and check of this lesson. It runs the full-batch training loop on the tickets and measures the result. Then it leaves out each step in turn, and shows the gradients that pile up without zero_grad(). Training takes a few seconds on a CPU. About 25 minutes.

Open in ColabOpen in KaggleView on GitHub

Colab and Kaggle open your own copy. They need a free Google or Kaggle account. Stuck? See the solution notebook.

Before you startWhat to know and to have readyAbout 50 minutes

What you need to know

• The loss and nn.CrossEntropyLoss(), from Loss: how wrong is the network?
• The update rule and the learning rate, from Gradient descent.
• loss.backward(), .grad, and gradients that add up, from Backpropagation and autograd.

What you need on your computer

• Nothing for the notebook route.
• For the local route: PyTorch and the folder ticket-router, from Set up PyTorch, and the setup lines of The tickets as numbers at the top of each code file.

Tested with Python 3.12 and PyTorch 2.14.1 on a CPU.

Practising a free throw

A basketball player practises free throws. Each throw has the same steps. She clears her head of the last correction. She throws. She sees how far the ball missed. She works out the correction: a little more power, a little to the left. Then she changes her throw, and starts again.

The order matters. If she changes her throw before she sees where the ball lands, she corrects for nothing. And if she does not clear her head, every old correction adds to the new one: after a few throws she corrects far too much.

The problem: the pieces must run in order, thousands of times

The last three lessons gave you the pieces. The loss says how wrong the network is, gradient descent changes the weights, and the backward pass gives the gradients. To train the ticket router, these pieces must run hundreds or thousands of times, always in the same order.

Each piece is one line of code, and a missing or misplaced line is easy to miss. Worse, PyTorch gives no error for most of these mistakes. The network trains, the loss is printed, and the result is just bad. So you need a fixed shape for training that you can check at a glance.

The idea: the five-step training loop

The training loop repeats five steps. One round of the five steps is one training step: one update of the weights.

The five steps of one training step, always in this order. The loop repeats them for the next batch of tickets.

Text description of the diagram

Five boxes in a cycle. Step 1, zero the gradients: optimizer.zero_grad(). Step 2, the forward pass: scores = model(x). Step 3, the loss: loss = loss_fn(scores, y). Step 4, the backward pass: loss.backward(). Step 5, update the weights: optimizer.step(). An arrow labelled “next batch” goes from step 5 back to step 1.

1. Zero the gradients. backward() adds to the gradients that are already there, so clear them first: the player clears her head.
2. Forward pass. The network computes its scores for the training tickets: the throw.
3. Loss. Compare the scores with the right teams: how far the ball missed.
4. Backward pass. A gradient for every parameter: the correction.
5. Update. Move every parameter against its gradient: change the throw.

Steps 1 and 5 belong to the optimizer: the object that holds the list of parameters and changes them. torch.optim.SGD(model.parameters(), lr=0.5) does exactly the update rule of the gradient descent lesson: every parameter minus 0.5 times its gradient. The course uses torch.optim.Adam with a learning rate of 0.001. Adam is also gradient descent, but it adapts the size of the step for each parameter. Learning rate and optimizers compares the two; for now, read Adam as “gradient descent that is harder to get wrong”.

Worked example: Priya trains the router

Priya starts with the simplest version. Every step uses all 4,696 training tickets at once. This is called full-batch training. She makes the optimizer once, before the loop, and logs the loss every 10 steps.

train_full_batch.py

set_seed(0)
model = BagOfWordsNet(len(vocab))
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for step in range(1, 101):
    optimizer.zero_grad()                  # 1. forget the last step's gradients
    scores = model(x_train)                # 2. forward pass: all 4,696 training tickets
    loss = loss_fn(scores, y_train)        # 3. how wrong?
    loss.backward()                        # 4. backward pass: a gradient for every parameter
    optimizer.step()                       # 5. update every parameter
    if step % 10 == 0:
        print(f"step {step:3}   loss {loss.item():.4f}")

The output:

step  10   loss 1.4113
step  20   loss 1.1342
step  30   loss 0.8550
step  40   loss 0.6181
step  50   loss 0.4454
step  60   loss 0.3313
step  70   loss 0.2573
step  80   loss 0.2076
step  90   loss 0.1724
step 100   loss 0.1460

model.parameters() gives the optimizer all 183,685 parameters of the network. step % 10 == 0 is true at every tenth step, so the log stays short. The whole run takes about 1.5 seconds on a CPU. The chart shows the loss of every step, not only every tenth.

The training loss at each of the 100 full-batch steps: Adam, learning rate 0.001, all 4,696 training tickets in every step. Real run, PyTorch 2.14.1 on a CPU.

Text description of the chart

The loss starts at 1.624, the loss of an even guess. It falls slowly for the first few steps, then quickly between steps 10 and 50 (from 1.41 to 0.45). Then it falls more slowly again: 0.26 at step 70 and 0.146 at step 100. It falls at every step; the curve has no jumps.

The loss falls at every step, with no jumps, unlike the plain gradient descent of the last lesson. Then Priya measures the trained network on the 600 validation tickets, which it never trained on:

train_full_batch.py

valid_loss, valid_acc = evaluate(model, x_valid, y_valid)
print(f"validation loss {valid_loss:.4f}   accuracy {valid_acc:.4f}")

The output:

validation loss 0.1985   accuracy 0.9783

From 18.2% to 97.8% in 100 steps. Omar asks what that means for Grace’s team: about 98 tickets in 100 now reach the right team without a person reading them. It is not yet better than the machine learning course’s TF-IDF model, which reaches 98.8% on the same tickets; the course compares them fairly at the end.

What happens when a step is missing

Priya trains for 20 steps four times: once with all five steps, and once without each of the first, fourth and fifth steps. Here is the loss at step 20 of each run, from the notebook:

Left out	Loss at step 20	Error message?	What happened
nothing	1.1342	no	The network learns.
zero_grad()	1.1745	no	It learns, but worse: old gradients pile up.
backward()	1.6236	no	No gradients, so step() has nothing to use. The network stays as it was.
step()	1.6236	no	The gradients are computed, then ignored. The network stays as it was.

None of the four runs stops with an error. A loss that does not move at all from its first value (1.6236 here) tells you that the weights never change: look for a missing backward() or step(). Leaving out step 2 or 3 is not possible: without them there is no loss to call backward() on.

When to use the loop, and what it costs

Every network that you train in PyTorch uses this loop. Higher-level libraries hide it inside a fit() or train() function, but the five steps are still there. The course’s ticketnet.py has one, train(), with the same five lines and the same comments; from the next lesson on, you can use it.

Full-batch training has two costs. First, memory: every step holds all the training tickets and the record of their forward pass. The word counts alone are 4,696 × 2,864 numbers, about 54 MB, and real datasets are millions of rows. Second, speed of learning: each step needs a full pass over the data, so 100 steps read every ticket 100 times. The next lesson fixes both with small batches.

Guided practice: plain gradient descent

1. In the notebook, go to section 4.
2. Make a new network with set_seed(0) and BagOfWordsNet(len(vocab)).
3. Make the optimizer: torch.optim.SGD(model.parameters(), lr=0.5).
4. Write the five steps inside a loop of 10 steps, and append each loss.item() to losses.
5. Run the check cell. You should see Check passed. The ten losses start at 1.6236 and end at 1.4191.

Show an example answer

train_sgd.py

losses = []
set_seed(0)
model = BagOfWordsNet(len(vocab))
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
for step in range(10):
    optimizer.zero_grad()
    loss = loss_fn(model(x_train), y_train)
    loss.backward()
    optimizer.step()
    losses.append(loss.item())

These are the same ten steps as the plain gradient descent on the ticket network in the last lesson: 1.624, 1.599, 1.578 … 1.419. torch.optim.SGD does the update rule for you. Adam reached 1.411 after the same ten steps, a little lower.

Your turn: a learning rate 10 times larger

Change Adam’s learning rate from 0.001 to 0.01 and train again for 100 full-batch steps. Before you run it, predict: will the training loss after 100 steps be lower or higher? And the validation accuracy?

Show the answer

Priya’s run: the training loss falls much faster, to 0.303 at step 10 and 0.0008 at step 100, almost zero. But the validation result is not better: the loss is 0.138 and the accuracy 96.7%, against 0.199 and 97.8% with 0.001. A training loss near zero means that the network has started to learn the training tickets by heart. The training loss alone cannot tell you how the network does on new tickets; the next lesson measures both, after every pass.

Common mistake

Forgetting zero_grad()

Symptom: Tomás deletes optimizer.zero_grad() “because the code still runs”. It does: no error, and the loss is printed. But the network learns worse and worse. Priya ran his version with the course’s usual small batches of 64 tickets, which the next lesson explains. After each of three passes through the data, the validation accuracy is 91.5%, 89.5% and 90.0%. The validation loss rises: 0.349, 0.417, 0.489. With zero_grad(), the same run gives 97.2%, 98.7% and 98.5%.

Cause: backward() adds each new gradient to the old ones. Without zero_grad(), every step is pushed by the sum of all the earlier gradients, which no longer fit the current weights. Priya measured the size of the first layer’s gradient after each pass: it grows from 10.7 to 15.7 to 25.6. With zero_grad(), it stays small: 0.66, 0.39, 0.39.

• With zero_grad()
• Without zero_grad()

The size (norm) of the first layer's gradient at the end of each pass, with and without zero_grad(). Batches of 64, Adam, learning rate 0.001. Real runs.

Text description of the chart

With zero_grad(), the gradient size is 0.66 after the first pass and 0.39 after the second and third. Without zero_grad(), it is 10.7, 15.7 and 25.6: it grows at every pass, because old gradients pile up.

Fix: call optimizer.zero_grad() as the first step of every training step, before the forward pass. To check a loop, print the gradient size of one layer, model.layers[0].weight.grad.norm(), for a few steps: it should not grow steadily.

What it means for a real AI service: a training mistake rarely crashes. It gives a model that is a little worse, and nobody notices unless the team checks the training log and the validation results of every run.

Knowledge check

Check your understanding

List the five steps of a training step in order, and say in a few words what each one does.

Show the answer

1, optimizer.zero_grad(): set the gradients to zero, because backward() adds to them. 2, the forward pass, model(x): compute the scores. 3, the loss, loss_fn(scores, y): how wrong the scores are. 4, loss.backward(): the backward pass, a gradient for every parameter. 5, optimizer.step(): move every parameter against its gradient.

Check your understanding

Predict: a training log shows the same loss, 1.6236, at every step. The code runs without errors. Which steps could be missing, and how do you check?

Show the answer

A loss that never moves means that the weights never change. So loss.backward() or optimizer.step() is missing, or the optimizer was made with the wrong parameters (for example, those of another model). Check that both lines are in the loop. Then print one parameter and its .grad before and after a step: the gradient should not be None, and the parameter should change.

Check your understanding

Tomás moves optimizer.zero_grad() to just after loss.backward(), before optimizer.step(). The code runs. What happens, and why?

Show the answer

The network stops learning. backward() computes the gradients, zero_grad() clears them at once, and step() then sees no gradients (zero or None), so the weights do not change. The loss stays at its first value. zero_grad() must come before the forward and backward passes of a step, not between backward() and step().

Recap

• A training step has five steps, in this order: zero_grad(), forward pass, loss, backward(), step(). The training loop repeats them.
• The optimizer holds the parameters and changes them: torch.optim.SGD is the update rule of gradient descent, torch.optim.Adam adapts the step of each parameter.
• 100 full-batch steps with Adam took the ticket network from a loss of 1.624 to 0.146 on the training tickets, and from 18.2% to 97.8% validation accuracy.
• A missing step rarely gives an error. A loss that never moves points to a missing backward() or step(); a missing zero_grad() makes the gradients pile up (10.7 → 15.7 → 25.6).
• Full-batch training reads every ticket in every step, which costs memory and time.

Next: train in small batches over several passes, and measure on the validation tickets after each pass.

Next: Batches, epochs and validation