Source: experiment.py

pip install torch numpy matplotlib
python3 experiment.py    # ~1 minute on a plain CPU; writes captures/ and diagrams/

#!/usr/bin/env python3
"""
Batch Norm vs Vanishing Gradients -- the one script behind the workshop.

Everything numeric in the workshop (every table, every plot, every quoted
number) comes from ONE deterministic run of this file:

    python3 experiment.py

Outputs:
    ../captures/results.log     -- every measurement, in citable sections
    ../diagrams/*.png           -- every matplotlib figure

Sections (match the [S*] markers in results.log):
    S1  two-moons dataset
    S2  sigmoid + BCE loss facts (max slope 0.25, ln2 baseline, dL/dz = p-y)
    S3  the networks (identical starting weights, verified)
    S4  at-init measurements: gradients, activations, slopes (plain vs BN)
    S5  geometric-decay illustration (pure arithmetic, no network)
    S6  training runs: plain vs BN
    S7  brute force vs depth: LR crank at depth 4, then depth 8
    S8  the other fixes: ReLU, Xavier init, residual, LayerNorm
    S9  the model.eval() bug, demonstrated

Requires: torch, numpy, matplotlib (CPU is plenty).
"""

import math
import os

import numpy as np
import torch
import torch.nn as nn
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# ----------------------------------------------------------------------------
# Setup: determinism and output paths
# ----------------------------------------------------------------------------

SEED = 0
DEPTH = 4          # hidden layers
WIDTH = 4          # neurons per hidden layer
N_SAMPLES = 1000   # two-moons points (500 per class)
NOISE = 0.10
N_TRAIN = 800
BATCH = 64
STEPS = 4000
LR = 0.5
EVAL_EVERY = 50

HERE = os.path.dirname(os.path.abspath(__file__))
CAPTURES = os.path.join(HERE, "..", "captures")
DIAGRAMS = os.path.join(HERE, "..", "diagrams")
os.makedirs(CAPTURES, exist_ok=True)
os.makedirs(DIAGRAMS, exist_ok=True)

LOG_PATH = os.path.join(CAPTURES, "results.log")
_log_file = open(LOG_PATH, "w")


def log(msg=""):
    """Print to stdout AND append to captures/results.log -- every number the
    workshop cites comes through this function."""
    print(msg)
    _log_file.write(msg + "\n")
    _log_file.flush()


def savefig(fig, name):
    """Save a figure under diagrams/ and note it in the citable log."""
    path = os.path.join(DIAGRAMS, name)
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    log(f"  [plot saved: diagrams/{name}]")


torch.manual_seed(SEED)
np.random.seed(SEED)

log("=" * 72)
log("Batch Norm vs Vanishing Gradients -- experiment log")
log(f"torch {torch.__version__} | numpy {np.__version__} | seed {SEED}")
log(f"network: {DEPTH} hidden layers x {WIDTH} neurons, sigmoid activations")
log("=" * 72)

# ----------------------------------------------------------------------------
# [S1] The dataset: two moons, hand-rolled (no sklearn needed)
# ----------------------------------------------------------------------------

log("\n[S1] two-moons dataset")

rng = np.random.default_rng(SEED)
n_half = N_SAMPLES // 2

# Upper moon: half circle of radius 1 centered at the origin.
t = rng.uniform(0.0, np.pi, size=n_half)
upper = np.stack([np.cos(t), np.sin(t)], axis=1)

# Lower moon: mirrored half circle, shifted to interleave with the upper one.
t = rng.uniform(0.0, np.pi, size=n_half)
lower = np.stack([1.0 - np.cos(t), 0.5 - np.sin(t)], axis=1)

# Stack both moons into one dataset and jitter every point with Gaussian noise.
X = np.concatenate([upper, lower]).astype(np.float32)
X += rng.normal(0.0, NOISE, size=X.shape).astype(np.float32)

# Labels: upper moon is class 0, lower moon is class 1.
y = np.concatenate([np.zeros(n_half), np.ones(n_half)]).astype(np.float32)

# Shuffle, then hold out the last 200 points as the never-trained-on test set.
perm = rng.permutation(N_SAMPLES)
X, y = X[perm], y[perm]
X_train, y_train = torch.from_numpy(X[:N_TRAIN]), torch.from_numpy(y[:N_TRAIN])
X_test, y_test = torch.from_numpy(X[N_TRAIN:]), torch.from_numpy(y[N_TRAIN:])

# Record the dataset's vital signs in the citable log.
log(f"  samples: {N_SAMPLES} ({n_half} per class), noise sigma = {NOISE}")
log(f"  split: {N_TRAIN} train / {N_SAMPLES - N_TRAIN} test")
log(f"  input feature std: x1 = {X[:,0].std():.3f}, x2 = {X[:,1].std():.3f}")

fig, ax = plt.subplots(figsize=(6.4, 4.2))
m0, m1 = y == 0, y == 1
ax.scatter(X[m0, 0], X[m0, 1], s=12, c="#1976d2", label="class 0")
ax.scatter(X[m1, 0], X[m1, 1], s=12, c="#d32f2f", label="class 1")
ax.set_xlabel("$x_1$")
ax.set_ylabel("$x_2$")
ax.set_title("The two-moons task: no straight line separates these")
ax.legend()
ax.grid(alpha=0.3)
savefig(fig, "two_moons.png")

# ----------------------------------------------------------------------------
# [S2] The two functions everything hinges on: sigmoid and BCE loss
# ----------------------------------------------------------------------------

log("\n[S2] sigmoid and loss facts")

# Evaluate the sigmoid on a fine grid.
zs = torch.linspace(-8, 8, 2001)
sig = torch.sigmoid(zs)

# Its slope has the closed form sigma'(z) = sigma(z) * (1 - sigma(z)).
dsig = sig * (1 - sig)

# Find the largest slope anywhere on the grid -- the claim is 0.25, at z = 0.
i = int(torch.argmax(dsig))
log(f"  max sigmoid slope     = {dsig[i]:.6f}  at z = {zs[i]:+.4f}")
log(f"  slope at z = +/-4     = {(torch.sigmoid(torch.tensor(4.0)) * (1 - torch.sigmoid(torch.tensor(4.0)))):.6f}")
log(f"  BCE loss at p = 0.5   = {-math.log(0.5):.6f}   (ln 2 = {math.log(2):.6f})")

# Verify dL/dz = p - y on a concrete example, letting autograd do the math:
# a logit of 2.0 scored against true label y = 0.
z = torch.tensor(2.0, requires_grad=True)
p = torch.sigmoid(z)
loss = nn.functional.binary_cross_entropy(p, torch.tensor(0.0))
loss.backward()
log(f"  autograd check at z = 2.0, y = 0:")
log(f"    p = sigmoid(2.0)    = {p.item():.6f}")
log(f"    formula p - y       = {p.item() - 0.0:+.6f}")
log(f"    autograd dL/dz      = {z.grad.item():+.6f}")

fig, ax = plt.subplots(figsize=(6.8, 4.0))
ax.plot(zs, sig, color="#1976d2", lw=2, label=r"$\sigma(z)$")
ax.plot(zs, dsig, color="#ef6c00", lw=2, label=r"slope $\sigma'(z)$")
ax.axhline(0.25, color="#ef6c00", ls=":", lw=1)
ax.annotate("max slope = 0.25", xy=(0, 0.25), xytext=(1.6, 0.36),
            arrowprops=dict(arrowstyle="->", color="#333"), fontsize=10)
ax.axvspan(-8, -4, color="grey", alpha=0.15)
ax.axvspan(4, 8, color="grey", alpha=0.15)
ax.text(-6, 0.8, "saturated\n(slope ~ 0)", ha="center", fontsize=9, color="#444")
ax.text(6, 0.8, "saturated\n(slope ~ 0)", ha="center", fontsize=9, color="#444")
ax.set_xlabel("z")
ax.set_title("The sigmoid and its slope")
ax.legend(loc="center left")
ax.grid(alpha=0.3)
savefig(fig, "sigmoid_and_deriv.png")

ps = torch.linspace(0.001, 0.999, 999)
fig, ax = plt.subplots(figsize=(6.8, 4.0))
ax.plot(ps, -torch.log(ps), color="#1976d2", lw=2, label="loss when y = 1:  $-\\ln p$")
ax.plot(ps, -torch.log(1 - ps), color="#d32f2f", lw=2, label="loss when y = 0:  $-\\ln(1-p)$")
ax.axhline(math.log(2), color="#555", ls=":", lw=1.2)
ax.annotate("ln 2 = 0.693 (the 'learned nothing' score,\nat p = 0.5)", xy=(0.5, math.log(2)),
            xytext=(0.53, 2.2), arrowprops=dict(arrowstyle="->", color="#333"), fontsize=9)
ax.set_xlabel("predicted probability p")
ax.set_ylabel("cross-entropy loss")
ax.set_ylim(0, 5)
ax.set_title("Binary cross-entropy loss")
ax.legend()
ax.grid(alpha=0.3)
savefig(fig, "bce_loss.png")

# ----------------------------------------------------------------------------
# [S3] The networks
# ----------------------------------------------------------------------------
# One class covers every variant in the workshop. `record=True` stores the
# per-layer pre-activations z_k and activations h_k and marks them so autograd
# keeps their gradients (retain_grad) -- that is how S4 measures the backward
# signal at every depth.


class TinyNet(nn.Module):
    """One network class for every variant in the workshop.

    plain / batch norm / LayerNorm come from `norm`; ReLU from `act`;
    residual from `residual`; the depth-8 nets from `depth`. Sharing one
    code path is itself a control: variants differ ONLY in these options.
    """

    def __init__(self, norm="none", act="sigmoid", residual=False, depth=DEPTH):
        super().__init__()
        # Set the manual seed so that different configurations (norm/act/...)
        # start with the exact same weight values, for a fair comparison.
        torch.manual_seed(SEED)

        # Dimensions: input (2 features) -> depth hidden layers of WIDTH -> output.
        dims = [2] + [WIDTH] * depth

        # One Linear layer per hidden level.
        self.hidden = nn.ModuleList(
            nn.Linear(dims[i], dims[i + 1]) for i in range(depth))

        # The normalizer that sits between each Linear and its activation.
        if norm == "batch":
            self.norms = nn.ModuleList(nn.BatchNorm1d(WIDTH) for _ in range(depth))
        elif norm == "layer":
            self.norms = nn.ModuleList(nn.LayerNorm(WIDTH) for _ in range(depth))
        else:
            # Identity() is a do-nothing placeholder: the plain net runs the
            # exact same forward() code path as the normalized variants.
            self.norms = nn.ModuleList(nn.Identity() for _ in range(depth))

        # Final layer projects to a single logit; the sigmoid lives in the loss.
        self.out = nn.Linear(WIDTH, 1)

        # Activation applied at every hidden layer.
        self.act = torch.sigmoid if act == "sigmoid" else torch.relu
        self.residual = residual
        self.rec = None   # intermediates from the last recorded forward pass

    def forward(self, x, record=False):
        # record=True stashes every layer's tensors and asks autograd to keep
        # their gradients (retain_grad) so S4 can read the backward signal at
        # every depth after loss.backward(). Off during training: it costs
        # memory and the optimizer only needs parameter gradients.
        rec = {"z_pre": [], "z": [], "h": []} if record else None
        h = x
        for lin, nrm in zip(self.hidden, self.norms):
            # The three steps of one hidden layer: linear, normalize, squash.
            z_pre = lin(h)
            z = nrm(z_pre)
            a = self.act(z)

            # Skip connection -- only Module 5's residual variant takes this.
            if self.residual and a.shape == h.shape:
                a = a + h

            if record:
                # Keep gradients on these non-leaf tensors for S4's readout.
                z.retain_grad()
                a.retain_grad()
                rec["z_pre"].append(z_pre)
                rec["z"].append(z)
                rec["h"].append(a)

            # This layer's activation feeds the next layer.
            h = a
        if record:
            self.rec = rec
        return self.out(h).squeeze(-1)


def weight_checksum(net):
    """Sum of |w| over all hidden-layer parameters: a cheap fingerprint used
    in S3 to PROVE two variants start from identical weights."""
    return sum(p.abs().sum().item() for lin in net.hidden for p in lin.parameters())


bce = nn.BCEWithLogitsLoss()   # sigmoid + binary cross-entropy, fused for stability


def accuracy(net, X, y):
    """Fraction of correct predictions (logit > 0 means p > 0.5).

    eval() first so batch norm normalizes with its running averages instead
    of THIS batch's statistics (the Module 5 bug, avoided); train() restores
    training behavior for the caller.
    """
    net.eval()
    with torch.no_grad():
        acc = ((net(X) > 0).float() == y).float().mean().item()
    net.train()
    return acc


def init_grad_ratio(net):
    """||dL/dW|| of the FIRST hidden layer over the output layer, before any
    training: the one-number fingerprint of vanishing gradients."""
    net.zero_grad()
    loss = bce(net(X_train), y_train)
    loss.backward()
    r = net.hidden[0].weight.grad.norm().item() / net.out.weight.grad.norm().item()
    net.zero_grad()
    return r


log("\n[S3] the networks")
plain = TinyNet(norm="none")
bn = TinyNet(norm="batch")
n_params = sum(p.numel() for p in plain.parameters())
log(f"  plain net parameters: {n_params}")
# The identical-start control: both variants must fingerprint identically.
log(f"  hidden-weight checksum, plain = {weight_checksum(plain):.6f}")
log(f"  hidden-weight checksum, bn    = {weight_checksum(bn):.6f}")
log(f"  identical starting weights: {weight_checksum(plain) == weight_checksum(bn)}")
# The starting weight scale: one of the two sub-unity factors in Module 2's
# per-layer hop product.
w1 = plain.hidden[1].weight
log(f"  layer-2 weight stats: std = {w1.std().item():.4f}, "
    f"max|w| = {w1.abs().max().item():.4f}  (PyTorch default init, fan_in = {WIDTH})")

# ----------------------------------------------------------------------------
# [S4] At-init measurements: one forward/backward pass, no training
# ----------------------------------------------------------------------------


def measure_at_init(net, label):
    """One forward/backward pass on the full training batch, no optimizer
    step: per layer, the weight gradient, the traveling signal dL/dh, the
    forward scales (std of z_pre / z / h), the mean sigmoid slope, the
    fraction of saturated units, and the typical gain W applies to a
    vector. This produces the S4 tables."""
    net.zero_grad()
    out = net(X_train, record=True)
    loss = bce(out, y_train)
    loss.backward()
    rec = net.rec
    depth = len(net.hidden)
    rows = []
    for k in range(depth):
        z, h, z_pre = rec["z"][k], rec["h"][k], rec["z_pre"][k]
        slopes = torch.sigmoid(z) * (1 - torch.sigmoid(z))
        slope = slopes.mean().item()

        # Saturation check beyond the mean: a healthy AVERAGE slope could
        # still hide a saturated tail, so also count the fraction of units
        # whose slope is below 0.05 (deep in the sigmoid's flat zones).
        sat_pct = (slopes < 0.05).float().mean().item() * 100

        # Typical gain a random vector receives through W_k: the
        # root-mean-square singular value = ||W||_F / sqrt(#singular values).
        # "Gain" is deliberately neutral: below 1 it shrinks, above 1 it
        # amplifies (some later configurations amplify).
        W = net.hidden[k].weight
        w_gain = W.norm().item() / math.sqrt(min(W.shape))
        rows.append(dict(
            layer=k + 1,
            w_grad=net.hidden[k].weight.grad.norm().item(),
            h_grad=h.grad.norm().item(),
            z_pre_std=z_pre.std().item(),
            z_std=z.std().item(),
            h_std=h.std().item(),
            slope=slope,
            sat_pct=sat_pct,
            w_gain=w_gain,
        ))
    out_grad = net.out.weight.grad.norm().item()
    log(f"\n  {label}: loss at init = {loss.item():.4f}")
    log(f"  {'layer':>7} {'||dL/dW||':>12} {'||dL/dh||':>12} {'std(z_pre)':>11} "
        f"{'std(z)':>8} {'std(h)':>8} {'mean slope':>11} {'%slope<.05':>11} {'W gain':>7}")
    for r in rows:
        log(f"  {r['layer']:>7} {r['w_grad']:>12.2e} {r['h_grad']:>12.2e} "
            f"{r['z_pre_std']:>11.4f} {r['z_std']:>8.4f} {r['h_std']:>8.4f} "
            f"{r['slope']:>11.4f} {r['sat_pct']:>10.1f}% {r['w_gain']:>7.3f}")
    log(f"  {'output':>7} {out_grad:>12.2e}")
    log(f"  gradient ratio, layer 1 vs output: "
        f"{rows[0]['w_grad'] / out_grad:.2e}  ({out_grad / rows[0]['w_grad']:.0f}x smaller)")
    factors = [rows[k]['h_grad'] / rows[k + 1]['h_grad'] for k in range(depth - 1)]
    log(f"  backward shrink factor per layer (||dL/dh_k|| / ||dL/dh_k+1||): "
        + ", ".join(f"{f:.3f}" for f in factors))
    net.zero_grad()
    return rows, out_grad


log("\n[S4] at-init measurements (full train batch, one backward pass)")
rows_plain, outg_plain = measure_at_init(plain, "PLAIN (Linear -> sigmoid)")
rows_bn, outg_bn = measure_at_init(bn, "BATCHNORM (Linear -> BN -> sigmoid)")

layers = list(range(1, DEPTH + 1))

fig, ax = plt.subplots(figsize=(6.8, 4.2))
ax.semilogy(layers + [DEPTH + 1], [r["w_grad"] for r in rows_plain] + [outg_plain],
            "o-", color="#d32f2f", lw=2, label="plain")
ax.semilogy(layers + [DEPTH + 1], [r["w_grad"] for r in rows_bn] + [outg_bn],
            "s-", color="#2e7d32", lw=2, label="with batch norm")
ax.set_xticks(layers + [DEPTH + 1])
ax.set_xticklabels([f"hidden {k}" for k in layers] + ["output"])
ax.set_ylabel(r"$\|\partial L/\partial W\|$  (log scale)")
ax.set_title("Weight-gradient size per layer, at initialization")
ax.legend()
ax.grid(alpha=0.3, which="both")
savefig(fig, "grad_norms_by_layer.png")

# The same chart with the plain network alone: Module 2 studies the plain
# net's cliff before batch norm has been introduced, so its copy of the
# figure must not spoil the (green) batch-norm line above.
fig, ax = plt.subplots(figsize=(6.8, 4.2))
ax.semilogy(layers + [DEPTH + 1], [r["w_grad"] for r in rows_plain] + [outg_plain],
            "o-", color="#d32f2f", lw=2, label="plain")
ax.set_xticks(layers + [DEPTH + 1])
ax.set_xticklabels([f"hidden {k}" for k in layers] + ["output"])
ax.set_ylabel(r"$\|\partial L/\partial W\|$  (log scale)")
ax.set_title("Weight-gradient size per layer at initialization: the plain net")
ax.legend()
ax.grid(alpha=0.3, which="both")
savefig(fig, "grad_norms_plain.png")

fig, ax = plt.subplots(figsize=(6.8, 4.2))
ax.semilogy(layers, [r["h_grad"] for r in rows_plain], "o-", color="#d32f2f", lw=2, label="plain")
ax.semilogy(layers, [r["h_grad"] for r in rows_bn], "s-", color="#2e7d32", lw=2, label="with batch norm")
ax.set_xticks(layers)
ax.set_xticklabels([f"$h_{k}$" for k in layers])
ax.set_ylabel(r"$\|\partial L/\partial h_k\|$  (log scale)")
ax.set_title("The backward signal at each depth, at initialization")
ax.legend()
ax.grid(alpha=0.3, which="both")
savefig(fig, "signal_decay.png")

fig, ax = plt.subplots(figsize=(6.8, 4.2))
ax.plot(layers, [r["h_std"] for r in rows_plain], "o-", color="#d32f2f", lw=2, label="plain: std of $h_k$")
ax.plot(layers, [r["h_std"] for r in rows_bn], "s-", color="#2e7d32", lw=2, label="batch norm: std of $h_k$")
ax.set_xticks(layers)
ax.set_xticklabels([f"$h_{k}$" for k in layers])
ax.set_ylabel("std of activations across the batch")
ax.set_title("Forward signal: how much the activations still vary, per layer")
ax.legend()
ax.grid(alpha=0.3)
savefig(fig, "act_std_by_layer.png")

# ----------------------------------------------------------------------------
# [S5] Geometric decay: what multiplying a small factor per layer does
# ----------------------------------------------------------------------------

log("\n[S5] geometric decay illustration (pure arithmetic)")
depths = np.arange(0, 11)
fig, ax = plt.subplots(figsize=(6.8, 4.2))
for f, c in [(0.9, "#2e7d32"), (0.5, "#ef6c00"), (0.25, "#d32f2f"), (0.1, "#6a1b9a")]:
    ax.semilogy(depths, f ** depths.astype(float), "o-", lw=2, color=c, label=f"factor {f} per layer")
    log(f"  factor {f}: after 4 layers -> {f**4:.2e}, after 10 layers -> {f**10:.2e}")
ax.set_xlabel("layers travelled backward")
ax.set_ylabel("remaining gradient size (log scale)")
ax.set_title("Multiplying one factor per layer is geometric decay")
ax.legend()
ax.grid(alpha=0.3, which="both")
savefig(fig, "geometric_decay.png")

# ----------------------------------------------------------------------------
# [S6] Training: plain vs batch norm
# ----------------------------------------------------------------------------


def train(net, lr=LR, steps=STEPS, label=""):
    """Plain minibatch SGD, deterministic (seeded sampler), recording the
    full-set training loss and held-out test accuracy every EVAL_EVERY
    steps. Every variant in the workshop trains through this one function."""
    opt = torch.optim.SGD(net.parameters(), lr=lr)

    # Seeded generator: the sequence of minibatches is identical across runs.
    g = torch.Generator().manual_seed(SEED)
    hist = {"step": [], "loss": [], "test_acc": []}
    net.train()
    for step in range(1, steps + 1):
        # Sample one minibatch and take one SGD step on its loss.
        idx = torch.randint(0, N_TRAIN, (BATCH,), generator=g)
        loss = bce(net(X_train[idx]), y_train[idx])
        opt.zero_grad()
        loss.backward()
        opt.step()

        # Periodically record full-set training loss and held-out accuracy.
        if step % EVAL_EVERY == 0 or step == 1:
            with torch.no_grad():
                net.eval()
                full_loss = bce(net(X_train), y_train).item()
                net.train()
            hist["step"].append(step)
            hist["loss"].append(full_loss)
            hist["test_acc"].append(accuracy(net, X_test, y_test))
    if label:
        first90 = next((s for s, a in zip(hist["step"], hist["test_acc"]) if a >= 0.9), None)
        first90 = f"step {first90}" if first90 else "never"
        log(f"  {label:<28} final train loss = {hist['loss'][-1]:.4f}, "
            f"final test acc = {hist['test_acc'][-1]*100:.1f}%, first >=90%: {first90}")
    return hist


log("\n[S6] training runs "
    f"(SGD, lr = {LR}, batch = {BATCH}, {STEPS} steps)")
plain = TinyNet(norm="none")
bn = TinyNet(norm="batch")
hist_plain = train(plain, label="plain")
hist_bn = train(bn, label="with batch norm")
log(f"  chance accuracy = 50.0% | 'learned nothing' loss = ln 2 = {math.log(2):.4f}")

fig, ax = plt.subplots(figsize=(6.8, 4.2))
ax.plot(hist_plain["step"], hist_plain["loss"], color="#d32f2f", lw=2, label="plain")
ax.plot(hist_bn["step"], hist_bn["loss"], color="#2e7d32", lw=2, label="with batch norm")
ax.axhline(math.log(2), color="#555", ls=":", lw=1.2)
ax.text(STEPS * 0.72, math.log(2) + 0.02, "ln 2 (learned nothing)", fontsize=9, color="#333")
ax.set_xlabel("training step")
ax.set_ylabel("training loss (full set)")
ax.set_title("Training loss: plain vs batch norm")
ax.legend()
ax.grid(alpha=0.3)
savefig(fig, "loss_curves.png")

fig, ax = plt.subplots(figsize=(6.8, 4.2))
ax.plot(hist_plain["step"], [a * 100 for a in hist_plain["test_acc"]], color="#d32f2f", lw=2, label="plain")
ax.plot(hist_bn["step"], [a * 100 for a in hist_bn["test_acc"]], color="#2e7d32", lw=2, label="with batch norm")
ax.axhline(50, color="#555", ls=":", lw=1.2)
ax.text(STEPS * 0.75, 51, "coin flip (50%)", fontsize=9, color="#333")
ax.set_xlabel("training step")
ax.set_ylabel("test accuracy (%)")
ax.set_ylim(40, 100)
ax.set_title("Test accuracy: plain vs batch norm")
ax.legend()
ax.grid(alpha=0.3)
savefig(fig, "accuracy_curves.png")


def decision_boundary(ax, net, title):
    """Evaluate the trained net's probability over a grid of the input plane
    and draw it, with the p = 0.5 contour as the decision boundary."""
    xs = np.linspace(-1.8, 2.8, 300)
    ys = np.linspace(-1.4, 1.9, 300)
    gx, gy = np.meshgrid(xs, ys)
    grid = torch.from_numpy(np.stack([gx.ravel(), gy.ravel()], axis=1).astype(np.float32))
    net.eval()
    with torch.no_grad():
        p = torch.sigmoid(net(grid)).reshape(gx.shape).numpy()
    ax.contourf(gx, gy, p, levels=20, cmap="RdBu_r", alpha=0.75, vmin=0, vmax=1)
    ax.contour(gx, gy, p, levels=[0.5], colors="k", linewidths=1.4)
    Xt, yt = X_test.numpy(), y_test.numpy()
    ax.scatter(Xt[yt == 0, 0], Xt[yt == 0, 1], s=9, c="#1976d2", edgecolors="none")
    ax.scatter(Xt[yt == 1, 0], Xt[yt == 1, 1], s=9, c="#d32f2f", edgecolors="none")
    ax.set_title(title, fontsize=11)
    ax.set_xticks([])
    ax.set_yticks([])


fig, axes = plt.subplots(1, 2, figsize=(9.6, 4.0))
decision_boundary(axes[0], plain, f"plain: {accuracy(plain, X_test, y_test)*100:.1f}% test accuracy")
decision_boundary(axes[1], bn, f"batch norm: {accuracy(bn, X_test, y_test)*100:.1f}% test accuracy")
fig.suptitle("What each trained network actually learned (black line = decision boundary)")
savefig(fig, "decision_boundaries.png")

# ----------------------------------------------------------------------------
# [S7] Brute force vs depth: crank the LR at depth 4, then go to depth 8
# ----------------------------------------------------------------------------

log(f"\n[S7] brute force vs depth ({STEPS} steps each)")
log(f"  depth {DEPTH} (the net from S6):")
crank_hists = {LR: hist_plain}
for lr in [5.0, 50.0]:
    net = TinyNet(norm="none")
    crank_hists[lr] = train(net, lr=lr, label=f"  plain, lr = {lr}")
log(f"    (lr = {LR} = the S6 run: final test acc = {hist_plain['test_acc'][-1]*100:.1f}%)")

fig, ax = plt.subplots(figsize=(6.8, 4.2))
for (lr, h), c in zip(crank_hists.items(), ["#d32f2f", "#ef6c00", "#6a1b9a"]):
    ax.plot(h["step"], [a * 100 for a in h["test_acc"]], lw=2, color=c, label=f"plain, lr = {lr}")
ax.axhline(50, color="#555", ls=":", lw=1.2)
ax.set_xlabel("training step")
ax.set_ylabel("test accuracy (%)")
ax.set_ylim(30, 100)
ax.set_title(f"Cranking the learning rate, plain net, depth {DEPTH}")
ax.legend()
ax.grid(alpha=0.3)
savefig(fig, "lr_crank.png")

DEEP = 8
log(f"\n  depth {DEEP} (same width, same recipe):")
p8, b8 = TinyNet(norm="none", depth=DEEP), TinyNet(norm="batch", depth=DEEP)
log(f"    init gradient ratio L1/out: plain = {init_grad_ratio(p8):.2e}, "
    f"batch norm = {init_grad_ratio(b8):.2e}")
deep_hists = {}
for lr in [0.5, 5.0, 50.0]:
    net = TinyNet(norm="none", depth=DEEP)
    deep_hists[f"plain, lr = {lr}"] = train(net, lr=lr, label=f"  plain, lr = {lr}")
deep_hists["batch norm, lr = 0.5"] = train(b8, label="  batch norm, lr = 0.5")

fig, ax = plt.subplots(figsize=(6.8, 4.2))
for (name, h), c in zip(deep_hists.items(),
                        ["#d32f2f", "#ef6c00", "#6a1b9a", "#2e7d32"]):
    ax.plot(h["step"], [a * 100 for a in h["test_acc"]], lw=2, color=c, label=name)
ax.axhline(50, color="#555", ls=":", lw=1.2)
ax.set_xlabel("training step")
ax.set_ylabel("test accuracy (%)")
ax.set_ylim(30, 100)
ax.set_title(f"Depth {DEEP}: no learning rate saves the plain net; batch norm just works")
ax.legend(fontsize=9)
ax.grid(alpha=0.3)
savefig(fig, "depth8_rescue.png")

# ----------------------------------------------------------------------------
# [S8] The other fixes
# ----------------------------------------------------------------------------

log(f"\n[S8] the other fixes (same init where applicable, same training recipe)")

variants = {}

# Xavier: choose the weight scale so each layer preserves signal variance,
# std = sqrt(2 / (fan_in + fan_out)) = 0.5 for our 4-into-4 layers.
xavier = TinyNet(norm="none")
for lin in list(xavier.hidden) + [xavier.out]:
    nn.init.xavier_uniform_(lin.weight)
    nn.init.zeros_(lin.bias)
log(f"  xavier layer-2 weight std = {xavier.hidden[1].weight.std().item():.4f} "
    f"(default was {TinyNet().hidden[1].weight.std().item():.4f})")

# The six contenders, all built from the same class with identical seeds.
candidates = {
    "plain sigmoid": TinyNet(norm="none"),
    "batch norm": TinyNet(norm="batch"),
    "ReLU": TinyNet(act="relu"),
    "Xavier init": xavier,
    "residual": TinyNet(residual=True),
    "layer norm": TinyNet(norm="layer"),
}
log(f"\n  {'variant':<16} {'init grad ratio L1/out':>23} {'final test acc':>15}")
for name, net in candidates.items():
    ratio = init_grad_ratio(net)
    hist = train(net)
    variants[name] = (ratio, hist["test_acc"][-1])
    log(f"  {name:<16} {ratio:>23.2e} {hist['test_acc'][-1]*100:>14.1f}%")

fig, ax = plt.subplots(figsize=(7.2, 4.0))
names = list(variants.keys())
accs = [variants[n][1] * 100 for n in names]
colors = ["#d32f2f" if a < 60 else "#2e7d32" for a in accs]
bars = ax.barh(names[::-1], accs[::-1], color=colors[::-1])
ax.axvline(50, color="#555", ls=":", lw=1.2)
ax.text(51, len(names) - 0.62, "coin flip", fontsize=9, color="#333")
ax.set_xlabel("final test accuracy (%)")
ax.set_xlim(0, 100)
ax.set_title("Every fix works; the plain sigmoid net does not")
for bar, a in zip(bars, accs[::-1]):
    ax.text(bar.get_width() + 1, bar.get_y() + bar.get_height() / 2, f"{a:.1f}%",
            va="center", fontsize=9)
ax.grid(alpha=0.3, axis="x")
savefig(fig, "variants_accuracy.png")

zs_r = torch.linspace(-4, 4, 1001)
relu = torch.relu(zs_r)
drelu = (zs_r > 0).float()
fig, ax = plt.subplots(figsize=(6.8, 4.0))
ax.plot(zs_r, relu, color="#1976d2", lw=2, label="ReLU(z)")
ax.plot(zs_r, drelu, color="#ef6c00", lw=2, label="slope of ReLU")
ax.annotate("slope exactly 1\n(no shrinkage)", xy=(2.5, 1.0), xytext=(0.7, 2.2),
            arrowprops=dict(arrowstyle="->", color="#333"), fontsize=10)
ax.set_xlabel("z")
ax.set_title("ReLU and its slope")
ax.legend(loc="upper left")
ax.grid(alpha=0.3)
savefig(fig, "relu_and_deriv.png")

# ----------------------------------------------------------------------------
# [S9] The classic bug: forgetting model.eval()
# ----------------------------------------------------------------------------

log("\n[S9] the model.eval() bug (using the trained batch-norm net from S6)")

bn.eval()    # correct: inference normalizes with the running averages
with torch.no_grad():
    acc_eval = ((bn(X_test) > 0).float() == y_test).float().mean().item()

bn.train()   # the bug: still normalizing every batch by its own statistics
with torch.no_grad():
    acc_train_full = ((bn(X_test) > 0).float() == y_test).float().mean().item()
    # batch-of-2 predictions in train mode (BN needs >1 sample to compute a std)
    preds = []
    for i in range(0, len(X_test) - 1, 2):
        preds.append((bn(X_test[i:i+2]) > 0).float())
    preds = torch.cat(preds)
    acc_train_b2 = (preds == y_test[:len(preds)]).float().mean().item()
bn.eval()

g0 = bn.norms[0].weight.detach()   # gamma, one per neuron
b0 = bn.norms[0].bias.detach()     # beta, one per neuron
log(f"  test accuracy, model.eval() (correct)          = {acc_eval*100:.1f}%")
log(f"  test accuracy, train mode, full test batch     = {acc_train_full*100:.1f}%")
log(f"  test accuracy, train mode, batches of 2        = {acc_train_b2*100:.1f}%")
log(f"  learned gamma (BN layer 1) = {[f'{v:.2f}' for v in g0.tolist()]}")
log(f"  learned beta  (BN layer 1) = {[f'{v:.2f}' for v in b0.tolist()]}")

log("\ndone. every figure written to diagrams/, every number above is citable.")
_log_file.close()