Multilayer Network — Programuoki

Code

The MLP — fill in two holes

In the last step you saw that two Linear layers with nothing between them are one Linear, and that ReLU breaks that. Now you write the same thing as a class — and fill in two missing pieces.

The network in words

An MLP (multi-layer perceptron) is this chain:

input → Linear → ReLU → Linear → ReLU → … → Linear → scores

hidden=(128,) means one hidden layer of 128 neurons: 784 → 128 → ReLU → 10.
hidden=(128, 64) — two: 784 → 128 → ReLU → 64 → ReLU → 10.
hidden=() — no hidden layer, the same as LinearClassifier.

The last layer has no ReLU. The output scores may be negative; ReLU would cut them off and CrossEntropyLoss would have nothing to compare. That is why the loop adds [Linear, ReLU] pairs, and the final Linear is appended separately.

nn.Sequential takes a list of layers and runs them one after another.

Your job

The model.py file has grown: LinearClassifier stays as it was (your file already has it), and an MLP class with a count_parameters function was added. Also create test_model.py — this time it is the one you pass.

Two holes:

HOLE 1 — a nn.ReLU() must follow every hidden Linear. Change the list so it holds both.
HOLE 2 — count_parameters must sum numbers, not tensors. The current line returns the number of tensors — it looks right but is wrong.

Run the test before filling anything in

$ python -m unittest test_model
FAIL: test_linear_classifier_has_7850_numbers
AssertionError: 2 != 7850
FAIL: test_mlp_128_has_101770_numbers
AssertionError: 4 != 101770
FAIL: test_two_hidden_layers
AssertionError: 6 != 109386
FAIL: test_doubling_the_input_does_not_double_the_output
AssertionError: True is not false
FAIL: test_learns_rings_which_a_single_layer_cannot
AssertionError: 0.5333333611488342 not greater than 0.95 : mlp reached 53.3%; linear got 53.3%
FAIL: test_one_relu_per_hidden_layer
AssertionError: 0 != 1 : hidden=(128,)
FAIL: test_no_hidden_layers_is_just_a_linear_classifier
AssertionError: 2 != 7850
Ran 11 tests in 1.2s

FAILED (failures=7)

(I kept only the FAIL and AssertionError lines; the real output also has a Traceback for each.)

That is the starting point. 2 != 7850: count_parameters counted two tensors (weight and bias) instead of 7 850 numbers. And without ReLU the network is linear: it does not learn the two rings (53.3 % — the same as a single layer).

Fill both holes and run again: OK, 11 tests.

Two numbers the test checks and that are worth understanding, not just passing:

784·128 + 128   =  100 480   hidden layer: weights + biases
128·10  + 10    =    1 290   output layer
                 ─────────
                  101 770

When it does not pass

TypeError: torch.nn.modules.activation.ReLU is not a Module subclass — you wrote nn.ReLU without parentheses. The list needs an object, not a class: nn.ReLU().
not greater than 0.95 : mlp reached 54.7% and 0 != 1 — you put in another activation (e.g. nn.Sigmoid()). The test insists on ReLU: it counts them. With sigmoid the network does learn, but not enough within those 400 steps.
20 != 7850, 276 != 101770 — you sum len(p), that is each tensor's first dimension. You need all elements: that is what p.numel() is for.
2 != 7850, but the ReLU tests pass — count_parameters is not changed yet.

model.py MODIFY 30 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


class MLP(nn.Module):
    """Linear -> ReLU -> Linear -> ReLU -> ... -> Linear."""

    def __init__(
        self,
        in_features: int = 784,
        hidden: tuple[int, ...] = (128,),
        num_classes: int = 10,
    ) -> None:
        super().__init__()
        layers: list[nn.Module] = []
        width = in_features
        for h in hidden:
            # HOLE 1: between two Linear layers goes an activation. Make the
            # list below hold the Linear AND, right after it, a ReLU.
            layers += [nn.Linear(width, h)]
            width = h
        layers.append(nn.Linear(width, num_classes))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


def count_parameters(model: nn.Module) -> int:
    # HOLE 2: count the NUMBERS the model holds, not the tensors it holds.
    # Replace the line below.
    return len(list(model.parameters()))

test_model.py CREATE

Create this file in ./

"""test_model.py — the verdict for model.py. You do not write this file."""

import unittest

import torch
from torch import nn

from model import MLP, LinearClassifier, count_parameters


def rings(n: int = 600) -> tuple[torch.Tensor, torch.Tensor]:
    """An inner ring (class 0, radius 0.5) inside an outer ring (class 1, radius 1)."""
    g = torch.Generator().manual_seed(0)
    y = torch.arange(n) % 2
    radius = 0.5 + 0.5 * y.float()
    angle = torch.rand(n, generator=g) * 6.2831853
    x = torch.stack([radius * angle.cos(), radius * angle.sin()], dim=1)
    return x + 0.05 * torch.randn(n, 2, generator=g), y


def train_full_batch(model: nn.Module, x, y, steps: int = 400, lr: float = 0.3) -> float:
    """Plain full-batch gradient descent, as in lesson 4. Returns train accuracy."""
    opt = torch.optim.SGD(model.parameters(), lr=lr)
    for _ in range(steps):
        loss = nn.functional.cross_entropy(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    with torch.no_grad():
        return (model(x).argmax(dim=1) == y).float().mean().item()


class TestCountParameters(unittest.TestCase):
    def test_linear_classifier_has_7850_numbers(self):
        self.assertEqual(count_parameters(LinearClassifier()), 7850)

    def test_mlp_128_has_101770_numbers(self):
        # 784*128 + 128   for the hidden layer
        # 128*10  + 10    for the output layer
        self.assertEqual(count_parameters(MLP()), 784 * 128 + 128 + 128 * 10 + 10)
        self.assertEqual(count_parameters(MLP()), 101_770)

    def test_two_hidden_layers(self):
        self.assertEqual(count_parameters(MLP(hidden=(128, 64))), 109_386)


class TestMLPShape(unittest.TestCase):
    def test_maps_784_to_10(self):
        out = MLP()(torch.zeros(5, 784))
        self.assertEqual(tuple(out.shape), (5, 10))

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

    def test_no_hidden_layers_is_just_a_linear_classifier(self):
        self.assertEqual(count_parameters(MLP(hidden=())), 7850)

    def test_layer_widths_follow_hidden(self):
        model = MLP(in_features=5, hidden=(16, 8), num_classes=3)
        shapes = [tuple(m.weight.shape) for m in model.modules() if isinstance(m, nn.Linear)]
        self.assertEqual(shapes, [(16, 5), (8, 16), (3, 8)])


class TestMLPIsNonLinear(unittest.TestCase):
    def test_one_relu_per_hidden_layer(self):
        for hidden in [(128,), (64, 32), (16, 16, 16)]:
            relus = [m for m in MLP(hidden=hidden).modules() if isinstance(m, nn.ReLU)]
            self.assertEqual(len(relus), len(hidden), f"hidden={hidden}")

    def test_no_activation_after_the_last_layer(self):
        last = list(MLP(hidden=(8,)).modules())[-1]
        self.assertIsInstance(last, nn.Linear)

    def test_doubling_the_input_does_not_double_the_output(self):
        # A purely linear map f satisfies f(2x) - f(0) = 2 * (f(x) - f(0)).
        torch.manual_seed(0)
        model = MLP(in_features=4, hidden=(32, 32), num_classes=3)
        x = torch.randn(10, 4)
        zero = model(torch.zeros(1, 4))
        lhs = model(2 * x) - zero
        rhs = 2 * (model(x) - zero)
        self.assertFalse(torch.allclose(lhs, rhs, atol=1e-5))

    def test_learns_rings_which_a_single_layer_cannot(self):
        x, y = rings()
        torch.manual_seed(1)
        linear = train_full_batch(LinearClassifier(2, 2), x, y)
        torch.manual_seed(1)
        mlp = train_full_batch(MLP(2, hidden=(16,), num_classes=2), x, y)
        self.assertLess(linear, 0.7, f"linear reached {linear:.1%}")
        self.assertGreater(mlp, 0.95, f"mlp reached {mlp:.1%}; linear got {linear:.1%}")


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