The Training Loop — Programuoki

Code

Run it and watch the number climb

When python -m unittest says OK, run the training itself:

$ python train.py

(The first time, load_fashion may download the data if you do not have it yet — you will see a long line of percentages. Silent afterwards.)

My run (CPU, ~4 s):

before training: train accuracy 8.5%
epoch  1  loss 0.7482  train accuracy 74.0%
epoch  2  loss 0.6696  train accuracy 75.3%
epoch  3  loss 0.5560  train accuracy 81.3%
epoch  4  loss 0.5384  train accuracy 81.5%
epoch  5  loss 0.5050  train accuracy 83.2%
epoch  6  loss 0.5305  train accuracy 81.7%
epoch  7  loss 0.5043  train accuracy 83.1%
epoch  8  loss 0.5402  train accuracy 80.7%
epoch  9  loss 0.5079  train accuracy 82.9%
epoch 10  loss 0.4405  train accuracy 85.1%
done in 4.0s

The seed is 42, so my run reproduces exactly. Yours will probably be close but not necessarily identical: other machines and PyTorch versions round differently.

What the numbers show

The untrained model guesses 8.5 %. That is not exactly 10 % — "chance" is a spread, not a number. In lesson 3 the untrained model scored 10.2 %; different random starting weights now gave 8.5 %.

Accuracy climbs, but not in a straight line. 74.0 → 75.3 → 81.3 → 81.5 → 83.2 → 81.7 → 83.1 → 80.7 → 82.9 → 85.1. Twice it dropped (epochs 6 and 8). Every update is made from 64 images — a noisy estimate of the direction — so the model "jitters" around the true downhill path. The trend is clear; any single epoch is not. So do not panic when you see 83.2 → 81.7.

This is not the same experiment as lesson 4. There it was 81.7 % after 300 steps; here 85.1 % after 1 570 updates (157 × 10). Also, nn.Linear picks its own starting weights, not your randn * 0.01. So 85.1 vs 81.7 does not show that "nn.Linear is better" — it shows the loop did more updates in the same time. That was the point: fast, without manual chores.

This number is training accuracy. On the very 10 000 images it learned from. How it looks on new ones — lesson 7.

The panel now shows the full version

In this step model.py and train.py are shown filled in. This is not asking you to rewrite them — it is the reference you can compare your files with. From now on the panel shows this version: later lessons change exactly it, and the changes will be highlighted as a difference from it.

If your files work and the tests pass but look different — good. What matters is that the tests pass and python train.py gives a similar number. Keep your version.

What did not change

Lesson 4's train_scratch.py stays where it was. Keep it: when a later layer or optimizer starts behaving strangely, the only way to see what optimizer.step() really does is to look at the file in which you wrote that line yourself.

model.py MODIFY 3 lines changed

Edit 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."""

    def __init__(self, in_features: int = 784, num_classes: int = 10) -> None:
        super().__init__()
        self.fc = nn.Linear(in_features, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc(x)  # raw scores (logits); the loss applies softmax itself

train.py MODIFY 16 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
            optimizer.zero_grad()  # clear old gradients
            loss.backward()  # backward
            optimizer.step()  # W -= lr * W.grad, for every parameter

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


def main() -> None:
    torch.manual_seed(42)
    x, y = load_fashion("train", n=10_000)
    loader = make_loader(x, y, batch_size=64, shuffle=True)

    model = LinearClassifier()
    start = time.time()
    _, acc = evaluate(model, loader)
    print(f"before training: train accuracy {acc:.1%}")

    fit(model, loader, epochs=10, lr=0.1)
    print(f"done in {time.time() - start:.1f}s")


if __name__ == "__main__":
    main()