A training loop by hand
Lesson 3 of 8
3A training loop by hand
The big pictureEvery training run, from linear regression to a language model with billions of parameters, is the same five lines: predict, compute the loss, zero the gradients, differentiate, step. Write them by hand here so no ready-made trainer is ever a black box again.
What the picture shows: The light circles the five stages in the same order on every step, from linear regression to a huge language model. The loss needs the prediction, backward needs the loss, and step needs the gradients.
Map of this part
The concepts in order, click any one to jump to it:
Two ways into the same idea
Training repeats four steps like a player practising shots: shoot the ball (forward pass), see how far it missed (the loss), understand what went wrong (backward pass), then adjust a little (optimiser step). Thousands of shots, each a little closer.
For each batch: pred = model(x), then loss = loss_fn(pred, y), then optimizer.zero_grad(), loss.backward() and optimizer.step(). You separate model.train() from model.eval(), and log training and validation loss every epoch to see whether it learns or memorises.
Before You Start
🔁Autograd
You know loss.backward() fills .grad for every weight and that gradients accumulate unless you zero them.
📈Linear regression
You know a simple model like y = wx + b and the mean squared error loss from the regression lesson.
The Core Idea
Every training run is the same five lines. In this lesson you write them by hand in NumPy, then see them in PyTorch, then train a network on a curve with them, then see what happens when you forget one.
1The loop by hand in NumPy
We start with a linear regression and no training library. The data is built from y = 3x − 1 plus noise, and the model starts at w = 0 and b = 0.
import numpy as np
rng = np.random.default_rng(0)
X = rng.uniform(-2, 2, size=(200, 1))
y = 3 * X[:, 0] - 1 + rng.normal(scale=0.2, size=200)
w, b = 0.0, 0.0
for epoch in range(6):
pred = w * X[:, 0] + b # forward
loss = ((pred - y) ** 2).mean() # loss
grad_w = 2 * ((pred - y) * X[:, 0]).mean() # backward
grad_b = 2 * (pred - y).mean()
w -= 0.3 * grad_w # step
b -= 0.3 * grad_b
print(f"epoch {epoch} loss={loss:.3f} w={w:.2f} b={b:.2f}")
The loss falls from 13.3 to 0.04, and w and b approach the true 3 and −1 we built the data with. It never reaches zero because we added noise with spread 0.2, whose square, 0.04, is the best possible.
Look at the picture: most of the learning happened in the first two epochs, then the steps became small. Here each epoch is one step, because we use all 200 points as one batch. With batches, an epoch is as many steps as there are batches.
2The same lines in PyTorch
Every line you wrote by hand has a PyTorch counterpart. The optimiser is what takes the step: SGD subtracts the gradient times the learning rate, and Adam adapts the step for each weight from its gradient history.
The model becomes an nn.Module instead of a formula, the loss a ready function, differentiation loss.backward() instead of the two gradient lines, and the step optimizer.step(). The order and the idea never change.
3A network learns a curve
Now a target no straight line can draw. Look at the picture: one input, a hidden layer of sixteen units followed by ReLU, and one output.
The target is y = x² plus noise, a curve no straight line can represent. One hidden layer of sixteen units with ReLU is enough.
import torch
from torch import nn
torch.manual_seed(0)
X = torch.linspace(-2, 2, 200).unsqueeze(1)
y = X ** 2 + 0.1 * torch.randn_like(X)
model = nn.Sequential(nn.Linear(1, 16), nn.ReLU(), nn.Linear(16, 1))
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.MSELoss()
for epoch in range(301):
pred = model(X)
loss = loss_fn(pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 100 == 0:
print(f"epoch {epoch:3} loss {loss.item():.4f}")
print("f(1.5) =", round(model(torch.tensor([[1.5]])).item(), 2), "true =", 1.5 ** 2)
The loss goes from 3.77 to 0.015 in three hundred epochs, and the prediction at 1.5 is 2.34 against the true 2.25. Note the five lines inside the loop: they are literally the same in every training script you will read from now on.
4Why zero_grad every step
backward adds to the previous gradients. Without zeroing, each weight moves by the sum of every gradient since the start. Look at the picture: same model, same data, and the difference is one line.
import torch
from torch import nn
torch.manual_seed(0)
X = torch.randn(64, 4)
y = X.sum(dim=1, keepdim=True)
for forget in [False, True]:
torch.manual_seed(0)
model = nn.Linear(4, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
for step in range(50):
loss = ((model(X) - y) ** 2).mean()
if not forget:
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"forgot zero_grad={forget}: final loss {loss.item():.4f}")
Same model, same data. With zeroing the loss is 0.0007; without it 1.98, roughly worse than where it started. No error, no warning, and the code "works". That is the most dangerous kind of bug: the silent one.
How It Works, Step by Step: how the weights move in one training step
The forward pass pred = model(xb) turns a batch of inputs into predictions and builds the graph.
The loss function compares predictions with targets and returns one number: MSELoss for regression and CrossEntropyLoss for classification.
optimizer.zero_grad() clears the old gradients, then loss.backward() computes a fresh gradient for every weight.
optimizer.step() reads .grad and moves each weight one step, then the next batch starts over.
Worked example
Start with the loop in plain NumPy, no training library, on a linear regression. Four stages per epoch, labelled by name.
softmax to the model output before CrossEntropyLoss. The function applies log_softmax internally, so you get softmax twice and the model learns far slower with no error. Always pass it the raw outputs.
ln(10) ≈ 2.30. If it is 15 or 0.01, something is wrong before training even starts.
Try it without code
Myth vs Reality
| Myth | Reality |
|---|---|
| Ready-made trainers do something smarter than these five lines. | Inside they run the same five lines. What they add is printing, saving and timing around the loop, not a different loop. |
| Forgetting zero_grad will raise an error you notice. | No error and no warning. In the lesson's example the loss ended at 1.98 instead of 0.0007, and the code "works". |
| The number of epochs is how many times the weights are updated. | Every batch is a step. 10,000 examples in batches of 100 for five epochs means 500 updates, not 5. |
Try it yourself
What is the correct order for one training step?
The loss needs the prediction, backward needs the loss, and step needs the gradients. Zeroing goes anywhere before backward. This order is fixed in every training script.
10,000 examples, batch size 100, 5 epochs. How many times is optimizer.step() called?
Each epoch has 10000 / 100 = 100 batches, one step each, so five epochs give 500 steps. The number of steps, not epochs, is what determines how many times the weights moved.
Quick review cards
zero_grad, backward, then optimizer.step(). Every training framework does this inside..grad and updates the weight. SGD subtracts the gradient times the learning rate; Adam adapts the step per weight from its gradient history.zero_grad every stepbackward adds to previous gradients. Without zeroing, the weights move by the sum of all gradients since the start, and training fails.Check Your Understanding
Why put optimizer.zero_grad() before loss.backward() rather than after optimizer.step()? Does it matter?
Both places work as long as zeroing happens once between each step and the next backward. Putting it right before backward is clearer because it guarantees clean gradients at the moment they are computed, even if someone later adds code between steps. What matters is that it never happens between backward and step, or you wipe the gradients before using them.
In the NumPy loop the loss stopped near 0.04 and never reached zero. Why?
Because the data was built with noise of spread 0.2, and its square, 0.04, is the lowest error possible. No model can predict random noise.
Resources for this lesson
Grouped by type. Start with the official docs for precision, or an article or video for the explanation.
Official docs6
Done with this lesson?
Sign in to save your progress.