The training loop

Now it really learns. All the parts are ready, and what holds them together is a loop of five lines: predict, measure the error, clear the gradients, run backwards, take a step. Then start over.

This is not a simplified teaching example. Those exact five lines sit in every PyTorch script in the world, only with different names in front of them. Anyone who has read someone else's training code recognises them instantly - and anyone who has typed them once reads someone else's training code differently from then on.

The goal

Now it really learns. And the insight is how little it takes: progress does not hide in one clever step, it hides in the repetition of a dumb one. Two hundred times a small step downhill, and a network that started out guessing hits the mark.

The order inside the loop is not arbitrary:

1. Prediction - builds this round's graph.
2. Loss - the one number at the top of the graph.
3. zero_grad() - clear away the gradients of the last round.
4. backward() - spread gradients from the top into every weight.
5. step() - push every weight a little downhill.

Step 3 has to sit between 2 and 4 (or right after 5, which is the same place one round later). Behind backward() it deletes exactly the gradients it just computed, and the network never moves. Missing entirely, the gradients of all previous rounds pile up and training derails after a few epochs.

The bridge: training loop and agent loop

This course lives on a platform about agents, and this chapter is the reason why. Put the two loops side by side:

Training: compute a prediction → measure the error → adjust the weights → start over.

Agent: ask the model → call a tool → hand back the result → start over.

Both times the core is a repetition, not a flash of genius. Both times a state is passed from round to round - the weights there, the message list here. Both times the actual work is not in the clever part but in the plumbing around it: data in, state forwarded, stopping condition. And both times what looks like a black box from the outside looks like a for block from the inside.

The tools and loop building block describes the second loop in detail. Once you have written the first one yourself, it reads like an old acquaintance - that is the same "understanding by building" that everything here is explained with, applied to a neighbouring subject.

Zero magic in this step

NumPy is allowed, torch, jax, tinygrad and autograd are forbidden. In this chapter you no longer import anything foreign anyway: everything the loop needs, you built yourself. That is the real test of the chapter.

The exercise

A new file, meintorch/train.py:

from .loss import mse_loss
from .optim import SGD


def train(model, x, y, epochs=100, lr=0.05, loss_fn=mse_loss):
    """Trains the model on the full dataset and returns the loss per epoch."""
    opt = SGD(model.parameters(), lr=lr)
    history = []
    for _ in range(epochs):
        # TODO(du): five lines - predict, measure the error, clear the
        # gradients, run backwards, take a step. Then append the loss as an
        # ordinary number to `history`.
        raise NotImplementedError
    return history

Hint 1 - the direction

The five lines are calls to things you already have: the model itself, loss_fn, opt.zero_grad(), loss.backward(), opt.step().

Watch the order of zero_grad() and backward(). If your loss does not move at all, those two are probably the wrong way round.

For history: loss.data is a NumPy scalar. It keeps the whole graph of the round alive as long as you hold on to it - float(...) cuts that tie and costs nothing.

Hint 2 - more concrete

        loss = loss_fn(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()
        history.append(float(loss.data))

If the loss falls and then tips into nan, lr is too large. Halve it.

def train(model, x, y, epochs=100, lr=0.05, loss_fn=mse_loss):
    opt = SGD(model.parameters(), lr=lr)
    history = []
    for _ in range(epochs):
        loss = loss_fn(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()
        history.append(float(loss.data))
    return history

In goal mode

The autonomous run has to prove that learning really happens - not that the function runs through:

Create `meintorch/train.py` with
`train(model, x, y, epochs=100, lr=0.05, loss_fn=mse_loss)`. The function
creates an SGD over `model.parameters()` and repeats per epoch: prediction,
loss, `zero_grad()`, `backward()`, `step()`. It returns the loss values as a
list of ordinary floats.
Mandatory: NumPy allowed, torch/jax/tinygrad/autograd forbidden. `zero_grad()`
sits before `backward()` so gradients do not pile up across epochs.
Write `tests/test_06_trainingsschleife.py`: on a regression task the loss has to
fall by at least a factor of 10 over 200 epochs, all values have to be finite, a
second call has to keep learning instead of derailing, and `epochs=0` must not
change any parameter.
Run `python -m pytest` until everything is green.

The milestone test

tests/test_06_trainingsschleife.py. The milestone is the first test: the loss at the end is at least ten times smaller than at the start. The third test is the one that exposes a forgotten zero_grad() - without cleaning up, a second run does not get better, it gets worse.

import numpy as np

from meintorch.loss import mse_loss
from meintorch.nn import Linear, ReLU, Sequential
from meintorch.tensor import Tensor
from meintorch.train import train


def task():
    # y = 2*x0 - x1 + 0.5, a task a small network learns reliably
    rng = np.random.default_rng(0)
    x = rng.normal(size=(64, 2))
    y = 2.0 * x[:, :1] - x[:, 1:] + 0.5
    net = Sequential(Linear(2, 8, seed=0), ReLU(), Linear(8, 1, seed=1))
    return net, Tensor(x), Tensor(y)


def test_the_loss_falls_across_the_epochs():
    net, x, y = task()
    history = train(net, x, y, epochs=200, lr=0.05)
    assert len(history) == 200
    assert all(np.isfinite(v) for v in history)
    assert history[-1] < history[0] / 10


def test_the_network_predicts_usefully_in_the_end():
    net, x, y = task()
    train(net, x, y, epochs=400, lr=0.05)
    assert float(mse_loss(net(x), y).data) < 0.05


def test_a_second_run_keeps_learning_instead_of_derailing():
    net, x, y = task()
    first = train(net, x, y, epochs=50, lr=0.05)
    second = train(net, x, y, epochs=50, lr=0.05)
    assert second[-1] < first[-1]


def test_zero_epochs_change_nothing():
    net, x, y = task()
    before = net.layers[0].W.data.copy()
    assert train(net, x, y, epochs=0) == []
    assert np.allclose(net.layers[0].W.data, before)