The Training Loop — Programuoki

Code

The training loop — fill in two holes

In the last step you saw the parts one at a time. Now they are put together into a training loop — and you fill in the two missing pieces.

The loop in words

For each epoch, for each batch:

1. loss      ← ask the model how wrong it was on this batch
2. clear       the gradients left over from the previous batch
3. backward    loss.backward()  — fills in the gradients
4. step        move every weight against its gradient

After each epoch the loop measures and prints the loss and the accuracy — the evaluate function for that is already written.

Notice the order: gradients are cleared before backward(), not after it. After backward() they are needed — the step is made from them.

Your job

Create train.py and test_train.py (see the panel on the right).

train.py compiles and runs, but learns nothing. The helpers make_loader and evaluate are already written (nothing new in them), and fit has two gaps marked HOLE 1 and HOLE 2:

HOLE 1 — clearing the gradients. The optimizer is already built in the line optimizer = torch.optim.SGD(...); you need one command of it.
HOLE 2 — the step itself. Again one optimizer command.

Replace each pass with a real line. Nothing else needs to change.

What does fit do when the model is not a LinearClassifier but a plain nn.Linear? Exactly the same. fit does not care what it is given as long as it is an nn.Module. That is why the test uses nn.Linear(2, 2) and two clouds of points — you write the model class in the next step.

Run the test before filling anything in

First as it is. In the project folder:

$ python -m unittest
.FF..
======================================================================
FAIL: test_accuracy_rises_on_separable_data (test_train.TestFit...)
----------------------------------------------------------------------
AssertionError: 0.5 not greater than 0.9 : accuracy 50.0% -> 50.0%

======================================================================
FAIL: test_one_epoch_is_exactly_the_manual_update_of_lesson_4 (test_train.TestFit...)
----------------------------------------------------------------------
AssertionError: False is not true : weight differs from the manual loop by up to 0.1795

----------------------------------------------------------------------
Ran 5 tests in 1.2s

FAILED (failures=2)

(The real output has a Traceback with file paths before each AssertionError — I kept only the essence.)

This is not an error but a starting point. The skeleton learns nothing: the accuracy stays at exactly 50.0 % (two classes, a coin toss), and the weights never move — so they differ from the manual loop. The make_loader and evaluate tests pass: those functions are already correct.

Fill both holes and run again — you should see OK.

How does the test know? The second test replays lesson 4's manual loop (p -= 0.1 * p.grad, p.grad.zero_()) from the same starting model over the same 4 batches, then compares the weights with what your fit produced. So it is not just "does accuracy rise" — it is "is this the same computation".

When it does not pass

weight differs from the manual loop by up to 0.2776, but the accuracy test passes — you skipped zero_grad() (or wrote just optimizer.zero_grad without parentheses, which does nothing and raises no error). It still learns, but the gradients pile up: every batch pushes the weights by the leftovers of old batches too.
accuracy 50.0% -> 50.0% and up to 0.1795 — no step was made. The usual cause: optimizer.step without () — Python raises nothing, it just does not run it.
both tests fail although step() and zero_grad() are both there — most likely swapped. A step before backward() would use gradients that are not filled in yet.

train.py CREATE

Create this file in ./

"""train.py — the training loop, automated: DataLoader, optimizer, nn.Module."""

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


def make_loader(
    x: torch.Tensor, y: torch.Tensor, batch_size: int = 64, shuffle: bool = False
) -> DataLoader:
    # num_workers=0 (the default) is the safe choice on Windows.
    return DataLoader(TensorDataset(x, y), batch_size=batch_size, shuffle=shuffle)


@torch.no_grad()
def evaluate(model: nn.Module, loader: DataLoader) -> tuple[float, float]:
    """Return (average loss, accuracy) over every batch of the loader."""
    loss_fn = nn.CrossEntropyLoss(reduction="sum")
    total_loss, correct, count = 0.0, 0, 0
    for xb, yb in loader:
        logits = model(xb)
        total_loss += loss_fn(logits, yb).item()
        correct += (logits.argmax(dim=1) == yb).sum().item()
        count += len(yb)
    return total_loss / count, correct / count


def fit(model: nn.Module, loader: DataLoader, epochs: int, lr: float) -> None:
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    for epoch in range(1, epochs + 1):
        for xb, yb in loader:
            loss = loss_fn(model(xb), yb)  # forward + loss
            # HOLE 1: clear the gradients left over from the previous batch.
            # Replace `pass` with the optimizer call that does it.
            pass
            loss.backward()  # backward
            # HOLE 2: move every parameter one step against its gradient.
            # Replace `pass` with the optimizer call that does it.
            pass

        train_loss, train_acc = evaluate(model, loader)
        print(f"epoch {epoch:2d}  loss {train_loss:.4f}  train accuracy {train_acc:.1%}")

test_train.py CREATE

Create this file in ./

"""test_train.py — the verdict for train.py. You do not write this file."""

import copy
import io
import unittest
from contextlib import redirect_stdout

import torch
import torch.nn.functional as F
from torch import nn

from train import evaluate, fit, make_loader


def two_blobs(n: int = 200) -> tuple[torch.Tensor, torch.Tensor]:
    """Two clouds of 2-D points, one around (-1, -1) and one around (+1, +1)."""
    g = torch.Generator().manual_seed(0)
    y = torch.arange(n) % 2
    centre = (y * 2 - 1).float().unsqueeze(1).repeat(1, 2)
    return centre + 0.5 * torch.randn(n, 2, generator=g), y


def run_fit(model: nn.Module, loader, epochs: int, lr: float) -> str:
    """Call fit() and return what it printed."""
    out = io.StringIO()
    with redirect_stdout(out):
        fit(model, loader, epochs, lr)
    return out.getvalue()


class TestMakeLoader(unittest.TestCase):
    def test_batches_cover_every_row_in_order(self):
        x = torch.arange(10).float().unsqueeze(1)
        y = torch.arange(10)
        loader = make_loader(x, y, batch_size=4, shuffle=False)
        sizes = [len(yb) for _, yb in loader]
        self.assertEqual(sizes, [4, 4, 2])
        self.assertEqual(torch.cat([yb for _, yb in loader]).tolist(), list(range(10)))


class TestEvaluate(unittest.TestCase):
    def test_loss_is_a_per_image_average_even_with_uneven_batches(self):
        model = nn.Linear(2, 2)
        with torch.no_grad():
            model.weight.copy_(torch.eye(2))
            model.bias.zero_()
        x = torch.tensor([[2.0, 0.0], [0.0, 3.0], [1.0, 0.0], [0.0, 0.5]])
        y = torch.tensor([0, 1, 1, 0])  # the scores are x itself: 2 of 4 are right
        loss, acc = evaluate(model, make_loader(x, y, batch_size=3))
        self.assertAlmostEqual(acc, 0.5)
        self.assertAlmostEqual(loss, F.cross_entropy(x, y).item(), places=5)


class TestFit(unittest.TestCase):
    def setUp(self):
        self.x, self.y = two_blobs()
        self.loader = make_loader(self.x, self.y, batch_size=50)  # 4 batches

    def test_one_epoch_is_exactly_the_manual_update_of_lesson_4(self):
        torch.manual_seed(1)
        by_fit = nn.Linear(2, 2)
        by_hand = copy.deepcopy(by_fit)

        run_fit(by_fit, self.loader, epochs=1, lr=0.1)

        for xb, yb in self.loader:  # lesson 4, once per batch
            loss = F.cross_entropy(by_hand(xb), yb)
            loss.backward()
            with torch.no_grad():
                for p in by_hand.parameters():
                    p -= 0.1 * p.grad
                    p.grad.zero_()

        for name, p in by_fit.named_parameters():
            q = dict(by_hand.named_parameters())[name]
            self.assertTrue(
                torch.allclose(p, q, atol=1e-6),
                f"{name} differs from the manual loop by up to "
                f"{(p - q).abs().max().item():.4f}",
            )

    def test_accuracy_rises_on_separable_data(self):
        torch.manual_seed(1)
        model = nn.Linear(2, 2)
        _, before = evaluate(model, self.loader)
        run_fit(model, self.loader, epochs=5, lr=0.1)
        _, after = evaluate(model, self.loader)
        self.assertGreater(after, 0.9, f"accuracy {before:.1%} -> {after:.1%}")

    def test_prints_one_line_per_epoch(self):
        torch.manual_seed(1)
        text = run_fit(nn.Linear(2, 2), self.loader, epochs=3, lr=0.1)
        lines = text.strip().splitlines()
        self.assertEqual(len(lines), 3)
        self.assertIn("train accuracy", lines[0])


if __name__ == "__main__":
    unittest.main()