Overfitting & Regularization

Code

Dropout, weight decay, early stopping — write it yourself

The split works. Now three measures against overfitting, and one rule without which they lie. This time with fewer hints.

Three measures

Dropout. While training, at every step a random part of the neurons is switched off (say 30 %), and the outputs of the rest are scaled up so the total is preserved. The network cannot rely on one neuron — next time it may not be there.

Weight decay. A penalty for large numbers is added to the loss: the bigger a weight, the more it costs. The network picks more modest, "smoother" functions. For Adam this is the parameter weight_decay.

Early stopping. We watch validation loss: if for patience epochs in a row it did not improve — we stop. And we bring back the best weights, not the ones we stopped at.

One rule: `train()` and `eval()`

Dropout has two modes: while training it switches neurons off, while measuring it does not. A model knows this from a single switch:

model.train()   # dropout is active
model.eval()    # dropout is off

What has to happen: fit turns train() on before every epoch, and evaluate turns eval() on before measuring. If neither — no error, no warning, but the numbers are wrong. In step 7 you will see by how much.

Your job

1. model.py — MLP gets a dropout parameter (default 0.0). After each hidden ReLU goes nn.Dropout(dropout); after the last layer — none. The change is small, so you see it finished in the panel; test_model.py checks it.

2. train.py:

evaluate — a one-line fix: turn on model.eval(). It is not in the panel.
fit — a skeleton with NotImplementedError; in the panel you see only the contract in comments. The change is large, so I do not offer "add two lines". Your old fit from lesson 6 is the starting point: keep the loop, add validation and stopping.
build_model and main are already updated (including --dropout, --weight-decay, --patience).

The fit contract (you will also find it in the comment):

for every epoch:
    model.train()
    ... the usual loop over batches ...
    train_loss, train_acc  ← evaluate(model, loader)
    if there is a val_loader:
        val_loss, val_acc  ← evaluate(model, val_loader)
        if val_loss is the lowest so far:
            remember it and a COPY of the weights;  waited ← 0;  row["best"] ← True
        else: waited ← waited + 1
    append the row to the history
    if patience is not given or waited < patience: continue
    else: stop

if patience is given and a copy exists: load the best weights back
return the history

Two hidden details:

the weights copy has to be a real copy (copy.deepcopy(model.state_dict())). state_dict() returns references to the live tensors — without deepcopy it simply changes along with the model;
the optimizer gets weight_decay: torch.optim.Adam(..., weight_decay=...).

3. The tests have grown. test_model.py got TestDropout, test_train.py got fit with validation and evaluate mode tests. Two changes in the old ones: build_model now takes a third argument (dropout), and the epoch line now says train acc, not train accuracy — so the validation numbers fit as well. The test was changed accordingly.

Run the tests

Before writing — almost everything fails with NotImplementedError, plus one real failure:

FAIL: test_evaluate_switches_dropout_off
AssertionError: ... same model, same data, two different answers

When everything is done: OK, 38 tests, ~1 s.

When it does not pass

test_dropout_is_on_while_learning_in_every_epoch fails, the list is ragged with False — you forgot model.train() at the start of the loop. Epoch one still works (a model starts in train() mode), and from epoch two on dropout is off — evaluate switched it off and nobody switched it back on.
... != ... within 4 places (test_patience_puts_the_best_weights_back) — you did not load the best weights back, or you saved a reference instead of a copy. Either way the model stays in the state of the last epoch.
not less than ... (test_weight_decay_shrinks_the_weights) — weight_decay did not reach the optimizer.
test_patience_stops_early_when_nothing_improves: len(history) is not 4 — check that waited is reset only on an improvement, not every epoch.

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


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

    def __init__(
        self,
        in_features: int = 784,
        hidden: tuple[int, ...] = (128,),
        num_classes: int = 10,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        layers: list[nn.Module] = []
        width = in_features
        for h in hidden:
            layers += [nn.Linear(width, h), nn.ReLU(), nn.Dropout(dropout)]
            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:
    return sum(p.numel() for p in model.parameters())

train.py MODIFY 53 lines changed

Edit this file in ./

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

import argparse
import copy
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,
    val_loader: DataLoader | None = None,
    weight_decay: float = 0.0,
    patience: int | None = None,
    verbose: bool = True,
) -> list[dict]:
    """Train; return one dict of metrics per epoch.

    patience: stop when val loss has not improved for this many epochs,
    and put the best weights back.
    """
    # Contract:
    #   - every epoch starts with model.train() (dropout ON while learning);
    #   - the optimizer is Adam, and it receives weight_decay;
    #   - after each epoch, measure the training set with evaluate() and append a
    #     dict: "epoch" (counting from 1), "train_loss", "train_acc";
    #   - if val_loader is given, also "val_loss" and "val_acc", and add
    #     "best": True to the row whenever val_loss is the lowest so far;
    #   - verbose=True prints one line per epoch in the format of lesson 5, with
    #     the validation numbers after a "|"; verbose=False prints nothing;
    #   - patience: once val_loss has failed to improve for `patience` epochs in
    #     a row, stop, print "early stop at epoch N: best val loss was X" (if
    #     verbose), and load the BEST weights back (keep a copy.deepcopy of
    #     model.state_dict() whenever val_loss improves).
    raise NotImplementedError


def build_model(name: str, hidden: tuple[int, ...], dropout: float) -> nn.Module:
    if name == "linear":
        return LinearClassifier()
    return MLP(hidden=hidden, dropout=dropout)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", choices=["linear", "mlp"], default="mlp")
    parser.add_argument("--hidden", type=int, nargs="+", default=[256, 256])
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--lr", type=float, default=0.001)
    parser.add_argument("--n-train", type=int, default=10_000)
    parser.add_argument("--dropout", type=float, default=0.0)
    parser.add_argument("--weight-decay", type=float, default=0.0)
    parser.add_argument("--patience", type=int, default=None)
    args = parser.parse_args()

    torch.manual_seed(42)
    x, y = load_fashion("train", n=args.n_train)
    xv, yv = load_fashion("val", n=5_000)
    loader = make_loader(x, y, batch_size=64, shuffle=True)
    val_loader = make_loader(xv, yv, batch_size=256)

    model = build_model(args.model, tuple(args.hidden), args.dropout)
    print(f"{args.model}: {count_parameters(model):,} parameters")
    start = time.time()
    history = fit(
        model,
        loader,
        epochs=args.epochs,
        lr=args.lr,
        val_loader=val_loader,
        weight_decay=args.weight_decay,
        patience=args.patience,
    )
    _, train_acc = evaluate(model, loader)
    _, val_acc = evaluate(model, val_loader)
    print(f"final: train {train_acc:.1%}  val {val_acc:.1%}  gap {train_acc - val_acc:+.1%}")
    print(f"done in {time.time() - start:.1f}s")


if __name__ == "__main__":
    main()

test_model.py MODIFY 30 lines changed

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


class TestDropout(unittest.TestCase):
    def test_default_is_off_and_holds_no_numbers(self):
        self.assertEqual(count_parameters(MLP(dropout=0.0)), 101_770)
        self.assertEqual(count_parameters(MLP(dropout=0.5)), 101_770)

    def test_one_dropout_per_hidden_layer_with_the_given_rate(self):
        model = MLP(hidden=(64, 32), dropout=0.3)
        drops = [m for m in model.modules() if isinstance(m, nn.Dropout)]
        self.assertEqual(len(drops), 2)
        self.assertTrue(all(d.p == 0.3 for d in drops))

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

    def test_dropout_changes_the_output_while_training(self):
        torch.manual_seed(0)
        model = MLP(in_features=4, hidden=(64,), num_classes=3, dropout=0.5)
        x = torch.randn(8, 4)
        model.train()
        self.assertFalse(torch.allclose(model(x), model(x)))

    def test_dropout_is_silent_in_eval_mode(self):
        torch.manual_seed(0)
        model = MLP(in_features=4, hidden=(64,), num_classes=3, dropout=0.5)
        x = torch.randn(8, 4)
        model.eval()
        self.assertTrue(torch.allclose(model(x), model(x)))


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

test_train.py MODIFY 125 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 noisy_blobs(n: int, seed: int) -> tuple[torch.Tensor, torch.Tensor]:
    """Two heavily overlapping clouds: easy to memorise when n is small."""
    g = torch.Generator().manual_seed(seed)
    y = torch.arange(n) % 2
    centre = (y * 2 - 1).float().unsqueeze(1).repeat(1, 2)
    return centre + 1.2 * torch.randn(n, 2, generator=g), y


class Spy(nn.Module):
    """A linear model that records whether it was in training mode on each call."""

    def __init__(self) -> None:
        super().__init__()
        self.fc = nn.Linear(2, 2)
        self.modes: list[bool] = []

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        self.modes.append(self.training)
        return self.fc(x)


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 acc", 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,), 0.0)
        self.assertIsInstance(model, LinearClassifier)
        self.assertEqual(count_parameters(model), 7850)

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

    def test_mlp_passes_dropout_on(self):
        model = build_model("mlp", (128,), 0.3)
        drops = [m for m in model.modules() if isinstance(m, nn.Dropout)]
        self.assertEqual([d.p for d in drops], [0.3])


class TestEvaluateMode(unittest.TestCase):
    def test_evaluate_switches_dropout_off(self):
        x, y = two_blobs()
        loader = make_loader(x, y, batch_size=200)
        model = nn.Sequential(nn.Linear(2, 2), nn.Dropout(0.9))
        model.train()
        first, second = evaluate(model, loader), evaluate(model, loader)
        self.assertEqual(first, second, "same model, same data, two different answers")
        self.assertFalse(model.training, "evaluate() left the model in training mode")


class TestFitWithValidation(unittest.TestCase):
    def setUp(self):
        self.x, self.y = noisy_blobs(30, seed=0)
        self.xv, self.yv = noisy_blobs(200, seed=1)
        self.loader = make_loader(self.x, self.y, batch_size=10, shuffle=True)
        self.val = make_loader(self.xv, self.yv, batch_size=200)

    def fit_quietly(self, model, **kw):
        out = io.StringIO()
        with redirect_stdout(out):
            history = fit(model, self.loader, val_loader=self.val, **kw)
        return history, out.getvalue()

    def test_one_row_per_epoch_with_train_and_val_numbers(self):
        torch.manual_seed(0)
        history, _ = self.fit_quietly(MLP(2, (8,), 2), epochs=4, lr=0.01)
        self.assertEqual([r["epoch"] for r in history], [1, 2, 3, 4])
        for key in ("train_loss", "train_acc", "val_loss", "val_acc"):
            self.assertTrue(all(key in r for r in history), f"missing {key}")

    def test_without_a_val_loader_there_are_no_val_keys(self):
        torch.manual_seed(0)
        out = io.StringIO()
        with redirect_stdout(out):
            history = fit(MLP(2, (8,), 2), self.loader, epochs=2, lr=0.01)
        self.assertEqual(len(history), 2)
        self.assertNotIn("val_loss", history[0])

    def test_best_marks_every_new_lowest_val_loss(self):
        torch.manual_seed(0)
        history, _ = self.fit_quietly(MLP(2, (8,), 2), epochs=6, lr=0.01)
        lowest = float("inf")
        for row in history:
            improved = row["val_loss"] < lowest
            self.assertEqual(row.get("best", False), improved, f"epoch {row['epoch']}")
            lowest = min(lowest, row["val_loss"])

    def test_verbose_false_prints_nothing_and_true_prints_a_line_per_epoch(self):
        torch.manual_seed(0)
        _, quiet = self.fit_quietly(MLP(2, (8,), 2), epochs=3, lr=0.01, verbose=False)
        self.assertEqual(quiet, "")
        _, loud = self.fit_quietly(MLP(2, (8,), 2), epochs=3, lr=0.01)
        lines = loud.strip().splitlines()
        self.assertEqual(len(lines), 3)
        self.assertIn("val acc", lines[0])

    def test_dropout_is_on_while_learning_in_every_epoch(self):
        spy = Spy()
        loader = make_loader(self.x, self.y, batch_size=10)  # 3 batches
        out = io.StringIO()
        with redirect_stdout(out):
            fit(spy, loader, epochs=2, lr=0.01)
        # per epoch: 3 training batches, then 3 measuring batches
        self.assertEqual(spy.modes, ([True] * 3 + [False] * 3) * 2)

    def test_patience_stops_early_when_nothing_improves(self):
        torch.manual_seed(0)
        # lr = 0: the weights never move, so val loss improves only in epoch 1.
        history, _ = self.fit_quietly(MLP(2, (8,), 2), epochs=50, lr=0.0, patience=3)
        self.assertEqual(len(history), 4)

    def test_patience_puts_the_best_weights_back(self):
        torch.manual_seed(0)
        model = MLP(2, (64, 64), 2)
        history, _ = self.fit_quietly(model, epochs=100, lr=0.01, patience=5)
        best = min(r["val_loss"] for r in history)
        self.assertLess(len(history), 100, "never stopped early")
        self.assertGreater(history[-1]["val_loss"], best + 0.05, "no overfitting to undo")
        restored, _ = evaluate(model, self.val)
        self.assertAlmostEqual(restored, best, places=4)

    def test_weight_decay_shrinks_the_weights(self):
        def total_size(weight_decay: float) -> float:
            torch.manual_seed(1)
            model = MLP(2, (16,), 2)
            out = io.StringIO()
            with redirect_stdout(out):
                fit(model, self.loader, epochs=30, lr=0.01, weight_decay=weight_decay)
            return sum(p.abs().sum().item() for p in model.parameters())

        self.assertLess(total_size(0.1), 0.8 * total_size(0.0))


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

Dropout, weight decay, early stopping — write it yourself

Three measures

One rule: train() and eval()

Your job

Run the tests

When it does not pass

One rule: `train()` and `eval()`