S101 · Regularization in DL — Dropout, BatchNorm, Weight Decay

"""regularization_lab.py — dropout + BN + weight decay, measured.
 
Trains four variants of the same MLP on MNIST:
  A. Vanilla       — no regularisation
  B. + Dropout
  C. + BatchNorm
  D. + Dropout + BN + AdamW weight_decay
Reports final train and test loss to show the generalisation gap.
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
 
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
class MLP(nn.Module):
    def __init__(self, use_dropout: bool = False, use_bn: bool = False,
                 dropout_p: float = 0.3) -> None:
        super().__init__()
        self.use_dropout = use_dropout
        self.use_bn = use_bn
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)
        if use_bn:
            self.bn1 = nn.BatchNorm1d(256)
            self.bn2 = nn.BatchNorm1d(128)
        if use_dropout:
            self.drop = nn.Dropout(p=dropout_p)
 
    def _block(self, x, linear, bn_layer):
        x = linear(x)
        if self.use_bn:
            x = bn_layer(x)
        x = F.relu(x)
        if self.use_dropout:
            x = self.drop(x)
        return x
 
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self._block(x, self.fc1, self.bn1 if self.use_bn else None)
        x = self._block(x, self.fc2, self.bn2 if self.use_bn else None)
        return self.fc3(x)
 
 
def make_loaders(batch: int = 128):
    tfm = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
    root = "./data"
    tr = datasets.MNIST(root, train=True,  download=True, transform=tfm)
    te = datasets.MNIST(root, train=False, download=True, transform=tfm)
    return DataLoader(tr, batch_size=batch, shuffle=True, num_workers=2), \
           DataLoader(te, batch_size=1024, shuffle=False, num_workers=2)
 
 
def run_epoch(model, loader, criterion, optimizer=None) -> tuple[float, float]:
    train = optimizer is not None
    model.train() if train else model.eval()
    total_loss, correct, n = 0.0, 0, 0
    with torch.set_grad_enabled(train):
        for x, y in loader:
            x, y = x.to(DEVICE), y.to(DEVICE)
            logits = model(x)
            loss = criterion(logits, y)
            if train:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
            total_loss += loss.item() * x.size(0)
            correct += (logits.argmax(1) == y).sum().item()
            n += x.size(0)
    return total_loss / n, correct / n
 
 
def train_variant(name: str, model: MLP, train_ld, test_ld,
                  *, epochs: int = 10, lr: float = 1e-3, wd: float = 0.0) -> None:
    model = model.to(DEVICE)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
    for ep in range(1, epochs + 1):
        tr_l, tr_a = run_epoch(model, train_ld, criterion, optimizer)
        te_l, te_a = run_epoch(model, test_ld,  criterion)
        if ep == epochs:
            gap = te_l - tr_l
            print(f"  [{name:<28s}] train_loss={tr_l:.4f} test_loss={te_l:.4f} "
                  f"test_acc={te_a:.4f} gap={gap:+.4f}")
 
 
if __name__ == "__main__":
    torch.manual_seed(0)
    train_ld, test_ld = make_loaders(batch=128)
 
    print("\nFinal-epoch metrics — LOWER gap = better generalisation:\n")
    train_variant("A · vanilla",              MLP(False, False), train_ld, test_ld, wd=0.0)
    train_variant("B · + dropout",            MLP(True,  False), train_ld, test_ld, wd=0.0)
    train_variant("C · + BN",                 MLP(False, True),  train_ld, test_ld, wd=0.0)
    train_variant("D · + dropout + BN + wd",  MLP(True,  True),  train_ld, test_ld, wd=0.05)