Multilayer Network — Programuoki

Code

The Adam optimizer and build_model — write it yourself

The model is ready. Now it has to be trained — and here plain SGD starts to fail for the first time.

Why we change the optimizer

The same MLP(128), the same slice of data, the same seed, 10 epochs, batch 64. Only the optimizer and lr change. All of this is my experiment, nothing from a "textbook":

sgd   lr 0.01   loss 0.6689  train acc 76.2%
sgd   lr 0.1    loss 0.9140  train acc 68.5%
sgd   lr 0.5    loss 0.7285  train acc 74.1%
adam  lr 0.001  loss 0.3622  train acc 86.3%
adam  lr 0.01   loss 0.3219  train acc 86.8%
adam  lr 0.1    loss 1.2670  train acc 45.6%

Three SGD lr values — and none comes close to what Adam gives on the first try. I am not saying SGD is "bad": I tried only three values and 10 epochs; with more epochs or momentum it would catch up. I am saying that here, with these settings, Adam hits first, and that its lr matters differently: 0.1 is already too much for it.

What Adam does: for every number separately it remembers how large and how consistent its latest gradients were, and picks that number's step accordingly. So:

the first step moves every number by roughly lr, whatever its gradient — the test checks exactly that;
numbers with noisy gradients move more carefully.

In your code this is one word: SGD → Adam. optimizer.zero_grad(), backward() and step() stay the same.

Your job

1. train.py: change torch.optim.SGD to torch.optim.Adam. Do not change anything else in fit. The file also grew in other ways: argparse (command-line parameters) and main() are provided — they just wire the parts together.

2. build_model(name, hidden) — an empty skeleton. The contract is in the docstring: "linear" returns LinearClassifier(), anything else MLP(hidden=hidden).

3. test_train.py: TestBuildModel tests were added and one test was replaced. Lesson 5's test that checked "exactly SGD's arithmetic" no longer applies — fit is no longer SGD. It was replaced by another: "after one step every number moved by about lr". That is not forgotten; it is a deliberate retirement.

Does fit in the panel have Adam? Yes — the change is already visible there. build_model still has NotImplementedError.

Run it

$ python -m unittest

While build_model is unwritten: two NotImplementedError errors. While fit still uses SGD, only one test fails:

FAIL: test_first_step_moves_every_number_by_about_lr_as_adam_does
AssertionError: False is not true : weight moved by [0.005079716444015503, 0.004353374242782593, 0.005079731345176697, 0.004353374242782593], want about 0.01 each

(SGD with lr=0.01 moves numbers not by 0.01 but by lr · gradient — here ~0.005.) When both are done: OK, 23 tests.

When it does not pass

AssertionError: 101770 != 109386 — build_model ignores hidden (you wrote MLP() instead of MLP(hidden=hidden)).
weight moved by [...], numbers not ~0.01 — the optimizer is not Adam (you probably left SGD).

train.py MODIFY 22 lines changed

Edit this file in ./

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

import argparse
import time

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

from fashion import load_fashion
from model import MLP, LinearClassifier, count_parameters


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.Adam(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 build_model(name: str, hidden: tuple[int, ...]) -> nn.Module:
    """Return a LinearClassifier for name == "linear", else an MLP.

    - "linear": a default LinearClassifier(); `hidden` is ignored;
    - anything else: MLP(hidden=hidden).
    """
    raise NotImplementedError


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", choices=["linear", "mlp"], default="mlp")
    parser.add_argument("--hidden", type=int, nargs="+", default=[128])
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--lr", type=float, default=0.001)
    args = parser.parse_args()

    torch.manual_seed(42)
    x, y = load_fashion("train", n=10_000)
    loader = make_loader(x, y, batch_size=64, shuffle=True)

    model = build_model(args.model, tuple(args.hidden))
    print(f"{args.model}: {count_parameters(model):,} parameters")
    start = time.time()
    _, acc = evaluate(model, loader)
    print(f"before training: train accuracy {acc:.1%}")

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


if __name__ == "__main__":
    main()

test_train.py MODIFY 25 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 MLP, LinearClassifier, count_parameters
from train import build_model, 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_first_step_moves_every_number_by_about_lr_as_adam_does(self):
        # Lesson 5 asserted SGD's arithmetic here. fit() now uses Adam, whose
        # very first step moves every number by roughly lr, whatever its gradient.
        torch.manual_seed(1)
        model = nn.Linear(2, 2)
        before = copy.deepcopy(model.state_dict())
        one_batch = make_loader(self.x, self.y, batch_size=len(self.y))

        run_fit(model, one_batch, epochs=1, lr=0.01)

        for name, p in model.named_parameters():
            moved = (p.detach() - before[name]).abs()
            self.assertTrue(
                torch.allclose(moved, torch.full_like(moved, 0.01), atol=1e-4),
                f"{name} moved by {moved.flatten().tolist()}, want about 0.01 each",
            )

    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)


class TestBuildModel(unittest.TestCase):
    def test_linear(self):
        model = build_model("linear", (128,))
        self.assertIsInstance(model, LinearClassifier)
        self.assertEqual(count_parameters(model), 7850)

    def test_mlp_uses_the_hidden_widths(self):
        self.assertIsInstance(build_model("mlp", (128,)), MLP)
        self.assertEqual(count_parameters(build_model("mlp", (128,))), 101_770)
        self.assertEqual(count_parameters(build_model("mlp", (128, 64))), 109_386)


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