Training loops, optimisers and learning rate

Monday’s support inbox had 31 new emails. Nineteen of them were about late deliveries. One customer had ordered a birthday cake for a 7pm party; it arrived at 8:40, after the candles had been blown out on a supermarket cake. The refund was ₹1,450, and the one-star review was free.

Meera forwarded it to Arjun with one line: “This is what the model is for.”

Arjun had the network from Friday: five inputs (distance, number of items, hour, rain, the delivery partner’s rating), two hidden layers, one output, 241 parameters. It computed its own mistakes perfectly. It just had not learned from any of them yet.

“How does it actually learn?” Meera asked when she came by. “Does it read the old orders?”

“Think of a new baker learning croissants,” Priya said. “Bake a tray. Taste them. Compare with the recipe. Work out what went wrong. Adjust. Bake again. A few thousand trays later, they’re good. Today Arjun writes the practice routine.”

The practice data

Crumb’s real delivery logs contain customer addresses, so they stay private. Instead, Arjun wrote a generator that produces realistic fake orders following rules he had observed. Save it as deliveries.py; the next lesson and the TensorFlow project reuse it.

# deliveries.py
import torch

FEATURES = ["distance_km", "items", "hour", "raining", "partner_rating"]

def make_deliveries(n: int = 2000, seed: int = 42):
    """Fake Crumb delivery data. Returns features X (n, 5) and labels y (n, 1)."""
    g = torch.Generator().manual_seed(seed)
    distance = torch.rand(n, generator=g) * 14 + 1             # 1 to 15 km
    items = torch.randint(1, 13, (n,), generator=g).float()     # 1 to 12 items
    hour = torch.randint(8, 23, (n,), generator=g).float()      # 8am to 10pm
    raining = (torch.rand(n, generator=g) < 0.3).float()        # rains 30% of the time
    rating = torch.rand(n, generator=g) * 2 + 3                 # partner rating 3.0 to 5.0
    rush = ((hour >= 18) & (hour <= 21)).float()                # evening rush hour

    # The hidden "true" rule. The model never sees this, only examples.
    logit = 0.3 * distance + 0.1 * items + 1.5 * raining + 1.5 * rush - 1.2 * (rating - 4) - 6
    late = (torch.rand(n, generator=g) < torch.sigmoid(logit)).float()

    X = torch.stack([distance, items, hour, raining, rating], dim=1)
    return X, late.unsqueeze(1)

About 20% of the generated orders are late (label 1). The rush-hour rule is the bump from last lesson that a single neuron cannot learn. The randomness in the last line matters: even with every feature known, some orders are late by bad luck, so no model can be perfect here, just as in real life.

Arjun’s first loop

Arjun had seen training loops in tutorials. He wrote one from memory.

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

X, y = make_deliveries(2000)
X = (X - X.mean(dim=0)) / X.std(dim=0)             # standardise, as in lesson 1

loader = DataLoader(TensorDataset(X, y), batch_size=64, shuffle=True)
model = nn.Sequential(nn.Linear(5, 16), nn.ReLU(), nn.Linear(16, 8), nn.ReLU(), nn.Linear(8, 1))
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

for epoch in range(20):
    running = 0.0
    for xb, yb in loader:
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()
        running += loss.item() * len(xb)
    print(f"epoch {epoch + 1:2d}  loss {running / len(X):.3f}")

The first few epochs looked great. Then it went strange.

epoch  1  loss 0.554
epoch  5  loss 0.427
epoch 10  loss 0.544
epoch 20  loss 0.569

The loss fell, then climbed back up. At the end, the model predicted “on time” for every single order. Since 80% of orders are on time, that scored 80% accuracy, and taught Meera nothing.

“Print the gradient size after each batch,” Priya suggested.

Arjun added print(model[4].weight.grad.norm()) inside the loop. On the first batch it was 0.12. After 10 batches, 0.61. After 200 batches, 1.94, sixteen times bigger than at the start, even though the model was supposedly getting better.

“Remember what the sous-chef does with the notebook?” Priya asked. “When you call backward, PyTorch adds the new gradient to whatever is already sitting in .grad. It never clears it by itself.”

“So every step was pushing with the sum of every gradient so far.”

“Like a baker who adjusts today’s recipe for every mistake they’ve made since January, all at once. The early mistakes were already fixed weeks ago. Now they’re being ‘fixed’ again, every single day.”

The fix is one line, optimizer.zero_grad(), before each backward. To understand where it goes, and why the other lines are there, follow one batch through all five steps.

Five steps, one tray of orders

Here is the loop again, correct this time. Read the numbered comments; this shape is the heart of every training script you will ever write, including the ones that train LLMs.

for epoch in range(20):
    running = 0.0
    for xb, yb in loader:                        # one mini-batch of 64 orders
        logits = model(xb)                       # 1. forward pass: predict
        loss = loss_fn(logits, yb)               # 2. measure how wrong
        optimizer.zero_grad()                    # 3. clear the old gradients
        loss.backward()                          # 4. backpropagate: compute new gradients
        optimizer.step()                         # 5. nudge every parameter
        running += loss.item() * len(xb)
    print(f"epoch {epoch + 1:2d}  loss {running / len(X):.3f}")

To see every number, Priya shrank the batch to four real orders from the generated data and ran one step with a freshly created model:

Order	Distance	Items	Hour	Rain	Partner rating	Actually late?
1	13.4 km	3	2pm	yes	3.86	no
2	13.8 km	8	1pm	no	3.80	yes
3	6.4 km	10	9pm	no	3.38	yes
4	14.4 km	12	noon	yes	4.75	yes

Step 1, forward pass. The untrained model’s raw scores (its logits) were 0.17, 0.15, 0.21 and 0.35. Through a sigmoid, that is 54%, 54%, 55% and 59% chance of being late. Essentially coin flips with a slight tilt, because the weights are random.

Step 2, loss. BCEWithLogitsLoss is binary cross-entropy (lesson 1’s “how wrong was this probability?” score for yes/no answers) with the sigmoid built in, averaged over the batch. Order 1 was on time but got 54% late: penalised. Orders 2 to 4 were late but only got 54 to 59%: also penalised, less. The average loss is 0.633. (A model that always says 50% scores 0.693, so this is barely better than a shrug.)

Step 3, zero_grad(). Every parameter’s .grad is reset to zero. On the first batch it is already empty; from the second batch onwards, this is what stops Arjun’s pile-up.

Step 4, backward(). Autograd fills in a gradient for all 241 parameters. The output layer’s eight weights got these slopes:

[ 0.0000, -0.0403, -0.0128,  0.0000, -0.0600, -0.0024,  0.0000, -0.0560]

Three are exactly zero. Those are weights attached to hidden neurons that ReLU had switched off for all four orders, so they could not have affected anything, just like the silent neuron in last lesson’s trace. The rest are negative, which means “increase me and the loss goes down”, because three of the four orders were late and the model should lean towards “late”.

Step 5, step(). Plain SGD applies the update rule from the gradient descent lesson to every parameter: new = old − 0.1 × gradient. The second output weight goes from 0.3167 to 0.3207. Tiny moves, everywhere at once.

Run the same four orders through again: the probabilities become 55%, 55%, 56% and 59%, and the loss drops from 0.633 to 0.626. Order 1, which was on time, also went up slightly. One step on four orders cannot please everyone; the next batch will pull in other directions. Over hundreds of batches, the pulls that agree with the real pattern add up, and the noise cancels out.

That is all training is. That loop, a few hundred or a few million times.

Batches, epochs and the DataLoader

A few words in that code deserve definitions.

The DataLoader takes your dataset and serves it in mini-batches (small random groups of examples). shuffle=True reshuffles the order every epoch, so each batch is a fair sample of the whole.

An epoch is one full pass through the training data. With 2,000 orders and batches of 64, one epoch is 32 batches (the last one has only 16 orders).

loss.item() pulls the plain Python number out of the loss tensor, so you can add it up for logging without dragging the whole computation graph along.

Inside step 4: backpropagation

loss.backward() runs backpropagation (computing every parameter’s gradient by applying the chain rule from the output back to the input, one layer at a time).

Priya told it as a complaint travelling backwards through the kitchen. A customer says the croissant was dense. The counter staff say it came out of the oven that way. The oven team say the dough was underproofed. The proofing team say the dough arrived cold. Each station passes the complaint back, scaled by how much it actually influenced the next station. A station that had no influence, like a hidden neuron that ReLU switched off, passes nothing back. At the end, every knob in the kitchen knows its share of the blame.

In the network, the loss depends on the output layer’s weights and on the second hidden layer’s values; those depend on their own weights and the first hidden layer’s values; and so on to the input. Backprop multiplies the local slopes along each path, which is exactly lesson 1’s chain rule, and the slopes of every path through a parameter are added up. It does all this in roughly the time of one extra forward pass, which is the clever bit: computing 241 gradients costs about as much as computing one prediction twice, not 241 times.

Inside step 5: the optimiser

The optimiser turns gradients into parameter changes. You hand it the parameters once, torch.optim.SGD(model.parameters(), lr=0.1), and each .step() applies its rule. The three you will meet most:

Plain SGD

The gradient descent rule, applied per mini-batch:

$$\theta \leftarrow \theta -\eta g$$

$\theta$ is the parameters, $\eta$ the learning rate, $g$ this batch’s gradient. Simple and predictable. Its weakness: each batch’s gradient is noisy, so it zigzags, and on long, gentle slopes it moves slowly because gentle slopes give small gradients.

SGD with momentum

Imagine pushing a loaded trolley of flour sacks across the storeroom. Once it is rolling, a bump in the floor does not knock it off course, and on a gentle downhill it keeps gathering speed. Momentum gives the update that same memory. Instead of following only this batch’s gradient, it follows a running blend of recent ones:

$$v\leftarrow \beta v+g\theta \leftarrow \theta -\eta v$$

$v$ is the “velocity”, and $\beta$ (usually 0.9) is how much of the previous direction to keep. Zigzags in opposite directions cancel out. Directions that keep agreeing build up.

Adam

Adam keeps momentum and tracks how big each parameter’s gradients usually are, then gives every parameter its own effective step size. A parameter with small, consistent gradients gets bigger steps; one with large, erratic gradients gets smaller ones. It is the default for most deep learning, usually with a learning rate around 1e-3. Its close cousin AdamW handles weight decay (a technique from next lesson) correctly and is the standard choice for transformers and LLM fine-tuning in Modules 30 and 33.

Here is how they compare on the delivery data, same model, same starting weights, 20 epochs:

Optimiser	Learning rate	Loss after epoch 5	Loss after epoch 20
SGD	0.01	0.598	0.471
SGD	0.1	0.417	0.370
SGD + momentum 0.9	0.01	0.434	0.370
Adam	0.001	0.469	0.368

Momentum at 0.01 catches plain SGD at 0.1, as the pause predicted. Adam gets there with a learning rate a hundred times smaller than that and no tuning. None of them gets much below 0.37, because of the bad-luck lateness built into the data: there is a floor no model can beat.

The learning rate, in practice

The learning rate is still the most important number in training, and the loss curve tells you whether it is right:

What the loss does	Likely cause	Try
Barely moves for many epochs	Too low	Multiply by 3 to 10
Drops, then bounces around wildly	Too high	Divide by 3 to 10
Becomes nan	Far too high, or bad input data	Divide by 10; check inputs for nan
Falls steadily, then flattens	Probably fine; maybe add a schedule	Try a scheduler

A learning rate scheduler changes the learning rate as training goes: big steps early to cover ground, small steps late to settle into the bottom of the valley. A common choice:

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)  # T_max = epochs

for epoch in range(20):
    for xb, yb in loader:
        ...                      # the five steps
    scheduler.step()             # once per epoch, after the batch loop

Large models also use a warm-up: a tiny learning rate for the first few hundred steps, rising to the full value, so the chaotic gradients of a random model do not throw it somewhere bad before it has learned anything.

Using the trained model

After training, switch to prediction mode. model.eval() tells layers that behave differently during training to switch modes (none in this model yet; next lesson adds one). torch.no_grad() switches off autograd’s notebook, saving memory and time.

model.eval()
with torch.no_grad():
    probs = torch.sigmoid(model(X))          # chance of late, 0 to 1
    preds = (probs > 0.5).float()            # 1 = "will be late"
    accuracy = (preds == y).float().mean()
print(f"accuracy {accuracy:.1%}")            # about 83.5%

83.5%. Arjun was ready to send a screenshot to Meera.

Where training goes wrong

Arjun’s missing zero_grad() is the most famous training-loop bug, but not the only one. Here are the others that silently ruin runs without crashing:

Double sigmoid. If the model ends with nn.Sigmoid() and you also use BCEWithLogitsLoss, the sigmoid is applied twice. Arjun tried it by accident a week later: the loss got stuck at 0.695 and every order was predicted “on time”. Either end with a raw nn.Linear and use BCEWithLogitsLoss (recommended, because it is numerically more stable), or end with a sigmoid and use BCELoss. Never both.

Unshuffled, sorted data. If the data happens to be sorted, say all on-time orders first, then an unshuffled loader feeds the model hundreds of “on time” batches, then hundreds of “late” ones. It lurches from one extreme to the other. In Arjun’s test it ended up flagging 35% of orders as late, when the truth was 20%, and accuracy fell to 75%. Use shuffle=True for training data.

Accumulating the loss tensor. running += loss (without .item()) keeps every batch’s computation graph alive. Long runs slowly eat all the memory.

A learning rate that worked yesterday. Change the batch size, model size or data, and the best learning rate moves. Re-check it.

Judging the model on the data it trained on. That 83.5% was measured on the same 2,000 orders the model learned from. That proves very little, as you are about to find out.

Try it yourself

Rerun the corrected loop four times, changing only the optimiser line. Put torch.manual_seed(0) before creating the model each time so they all start from the same random weights:

1. torch.optim.SGD(model.parameters(), lr=0.01)
2. torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
3. torch.optim.Adam(model.parameters(), lr=1e-3)
4. torch.optim.SGD(model.parameters(), lr=5.0)

Record the loss at epochs 5 and 20.

Show the answer

Runs 1 to 3 should roughly match the table above (your exact numbers will differ a little with different PyTorch versions and shuffles): plain SGD at 0.01 is clearly slowest; momentum and Adam both reach about 0.37.

Run 4 is the instructive one. With a learning rate of 5, the loss usually jumps around, gets stuck high, or becomes nan. Often the ReLUs “die”: one huge step throws the weights so far that every hidden neuron outputs zero for every input. Then every gradient is zero, as in the four-order trace above but for all neurons, and the model stops learning entirely. You will meet dead ReLUs properly in Module 28.

Common mistakes

• Forgetting optimizer.zero_grad(). Gradients pile up and updates become nonsense, often slowly enough that you do not notice for a while. Write the five steps in the same order every time until it is muscle memory.
• Applying sigmoid before BCEWithLogitsLoss. The loss already applies it. Output raw logits.
• Accumulating loss instead of loss.item(). A slow memory leak that crashes long runs.
• Not shuffling training data. Batches stop being fair samples, and the model lurches between extremes.
• Measuring success on training data. Fixed in the next lesson.

Recap

• Every training loop repeats five steps per mini-batch: forward, loss, zero_grad(), backward(), step().
• An epoch is one pass over the data; the DataLoader serves shuffled mini-batches; steps = batches per epoch × epochs.
• Backpropagation applies the chain rule backwards through the network and costs about one extra forward pass.
• SGD follows the gradient; momentum adds a rolling memory (up to 10× the step on steady slopes with β = 0.9); Adam adds a per-parameter step size and is the usual default.
• The learning rate matters more than anything else; schedulers change it as training goes.

Arjun showed Meera the 83.5% and explained the loop in her terms. “The model guesses for 64 orders at a time. We score how wrong it was. PyTorch works out which of the 241 numbers to blame, and we move each one a tiny bit. We wipe the blame clean and do the next 64. After 640 rounds, it’s right 83.5% of the time.”

Meera was delighted. Priya was not. “Arjun, what fraction of orders are late?”

“About 20%.”

“So a model that answers ‘on time’ for every order, without looking at anything, scores…?”

Arjun’s face fell. 80%. His 241-parameter model was 3.5 points better than a rubber stamp, measured on the very orders it had practised on. Tomorrow: how to tell a good model from a lucky one.

Go deeper

• PyTorch tutorials: the “Optimizing Model Parameters” page in “Learn the Basics” walks through this same loop.
• If you want to build autograd yourself in about 100 lines of Python, Andrej Karpathy’s “spelled-out intro to neural networks and backpropagation” is the best two hours you can spend this month.