The Training Loop — Programuoki

Code

The model as an nn.Module — write it yourself

The training loop works with any nn.Module. Now you write your own model and the function that wires everything together. This time with no gaps: empty skeletons.

fit in the panel still shows pass where your two lines are — the panel shows the skeleton, not your file. In your copy they are already filled in; leave them.

`nn.Module` — how it is written

A model is a class that inherits from nn.Module and has two methods:

class Something(nn.Module):
    def __init__(self, ...):
        super().__init__()          # first — always
        self.layer = nn.Linear(...) # layers live as attributes

    def forward(self, x):
        return ...                  # what the model does with its input

Three rules worth knowing before you write:

super().__init__() is the first line. Without it Python raises AttributeError: cannot assign module before Module.__init__() call.
Assign the layer to self.something = .... Only then does it show up in model.parameters() — and so in the optimizer. A layer that lives only in a local variable will not be trained.
forward must return. Without it the model returns None and the loss complains: argument 'input' (position 1) must be Tensor, not NoneType.

You call a model like a function: model(x) calls forward(x) inside.

Why the model does not return probabilities

In lesson 4 forward ended with softmax. Here it does not. The model returns raw scores (logits) and nn.CrossEntropyLoss applies softmax itself, inside. If you applied softmax twice there would be no error — just somewhat worse learning. You will see the numbers in the next step, "When the loop lies".

Your job

1. model.py (a new file) — the class LinearClassifier. The contract is in the comments. It must be the same 10 neurons from lesson 4, now as one object: scores = x @ W + b.

in_features pixels in, num_classes scores out;
one nn.Linear inside — the default model owns exactly 7 850 numbers;
forward returns raw scores.

2. train.py has grown: import lines and main() — an empty skeleton. Its docstring spells out exactly what to do and in what order (seed 42, 10 000 images, batch 64, 10 epochs, lr=0.1).

main has no unit test — its "test" is the next step: you run it and see whether you get what lesson 4 led you to expect.

3. test_train.py has grown: a TestLinearClassifier class. It checks the shape [N, 784] → [N, 10], 7 850 numbers, that it is an nn.Module, that raw scores come back (not probabilities), other sizes (LinearClassifier(2, 3)), and that fit trains this model.

Run python -m unittest. While you write, the test tells you what is wrong — the same kind of messages as in the previous step. Until the model is written you will see NotImplementedError.

model.py CREATE

Create this file in ./

"""model.py — the model as an nn.Module."""

import torch
from torch import nn


class LinearClassifier(nn.Module):
    """The same 10 neurons as lesson 4: scores = x @ W + b."""

    # Contract:
    #   - in_features pixels in, num_classes scores out: [N, in_features] -> [N, num_classes];
    #   - ONE nn.Linear inside, so a default model owns exactly 7 850 numbers;
    #   - forward returns the raw scores ("logits"). No softmax here: the loss
    #     function applies it itself.

    def __init__(self, in_features: int = 784, num_classes: int = 10) -> None:
        raise NotImplementedError

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        raise NotImplementedError

train.py MODIFY 19 lines changed

Edit this file in ./

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

import time

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

from fashion import load_fashion
from model import LinearClassifier


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%}")


def main() -> None:
    """Train the linear classifier on 10 000 FashionMNIST images and report.

    In this order:
      1. torch.manual_seed(42), then load_fashion("train", n=10_000);
      2. a loader with batch_size=64, shuffle=True;
      3. a LinearClassifier; print the accuracy BEFORE training as
         "before training: train accuracy 8.5%";
      4. fit(...) for 10 epochs with lr=0.1;
      5. print how long the whole run took: "done in 5.2s".
    """
    raise NotImplementedError

test_train.py MODIFY 31 lines changed

Edit 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 model import LinearClassifier
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])


class TestLinearClassifier(unittest.TestCase):
    def test_maps_784_pixels_to_10_scores(self):
        out = LinearClassifier()(torch.zeros(5, 784))
        self.assertEqual(tuple(out.shape), (5, 10))

    def test_is_an_nn_module_with_lesson_4s_7850_numbers(self):
        model = LinearClassifier()
        self.assertIsInstance(model, nn.Module)
        self.assertEqual(sum(p.numel() for p in model.parameters()), 7850)

    def test_other_sizes(self):
        out = LinearClassifier(in_features=2, num_classes=3)(torch.zeros(4, 2))
        self.assertEqual(tuple(out.shape), (4, 3))

    def test_returns_raw_scores_not_probabilities(self):
        torch.manual_seed(0)
        out = LinearClassifier()(torch.rand(8, 784))
        self.assertTrue((out < 0).any(), "no negative score: did you apply softmax?")
        self.assertFalse(torch.allclose(out.sum(dim=1), torch.ones(8)))

    def test_learns_with_fit(self):
        x, y = two_blobs()
        loader = make_loader(x, y, batch_size=50)
        torch.manual_seed(1)
        model = LinearClassifier(in_features=2, num_classes=2)
        run_fit(model, loader, epochs=5, lr=0.1)
        _, acc = evaluate(model, loader)
        self.assertGreater(acc, 0.9)


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

The model as an nn.Module — write it yourself

nn.Module — how it is written

Why the model does not return probabilities

Your job

`nn.Module` — how it is written