S097 · Multi-Layer Perceptron — Forward Pass

"""mlp_lab.py — a 2-hidden-layer MLP on MNIST, pure numpy.
 
Focus: THE FORWARD PASS AS MATRIX OPS.
Training uses a hand-rolled SGD with hand-derived gradients so you can see
the whole thing in one file. Session 098 will formalise backprop.
"""
from __future__ import annotations
import numpy as np
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
 
RNG = np.random.default_rng(0)
 
 
# ---------------- activations ----------------
def relu(z):        return np.maximum(0, z)
def relu_grad(z):   return (z > 0).astype(z.dtype)
def softmax(z):
    z = z - z.max(axis=1, keepdims=True)       # numerical stability
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)
 
 
# ---------------- MLP ----------------
class MLP:
    """Input → H1 → H2 → num_classes, ReLU hidden, softmax output."""
 
    def __init__(self, layer_sizes: list[int]) -> None:
        # He initialisation — right for ReLU
        self.W: list[np.ndarray] = []
        self.b: list[np.ndarray] = []
        for i in range(len(layer_sizes) - 1):
            fan_in, fan_out = layer_sizes[i], layer_sizes[i + 1]
            self.W.append(RNG.normal(0, np.sqrt(2.0 / fan_in), (fan_in, fan_out)))
            self.b.append(np.zeros(fan_out))
 
    def forward(self, X: np.ndarray) -> tuple[np.ndarray, dict]:
        """Return probabilities and cache the intermediate values for backward."""
        a = X
        cache = {"a": [X], "z": []}
        # hidden layers with ReLU
        for i in range(len(self.W) - 1):
            z = a @ self.W[i] + self.b[i]
            a = relu(z)
            cache["z"].append(z)
            cache["a"].append(a)
        # output layer with softmax
        z_out = a @ self.W[-1] + self.b[-1]
        probs = softmax(z_out)
        cache["z"].append(z_out)
        cache["a"].append(probs)
        return probs, cache
 
    def loss_and_grads(self, probs, y_onehot, cache):
        """Cross-entropy loss + hand-derived gradients."""
        n = y_onehot.shape[0]
        loss = -np.mean(np.sum(y_onehot * np.log(probs + 1e-12), axis=1))
 
        dW, db = [None] * len(self.W), [None] * len(self.b)
        dz = (probs - y_onehot) / n     # dL/dz for softmax + CE
        for i in reversed(range(len(self.W))):
            a_prev = cache["a"][i]
            dW[i] = a_prev.T @ dz
            db[i] = dz.sum(axis=0)
            if i > 0:
                da_prev = dz @ self.W[i].T
                dz = da_prev * relu_grad(cache["z"][i - 1])
        return loss, dW, db
 
    def step(self, dW, db, lr: float) -> None:
        for i in range(len(self.W)):
            self.W[i] -= lr * dW[i]
            self.b[i] -= lr * db[i]
 
 
# ---------------- data ----------------
def load_mnist():
    print("Downloading MNIST (may take 30-60s the first time)...")
    ds = fetch_openml("mnist_784", version=1, as_frame=False, parser="auto")
    X = (ds.data.astype(np.float32) / 255.0)
    y = ds.target.astype(int)
    Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=10_000, random_state=0)
    return Xtr[:20_000], ytr[:20_000], Xte, yte    # subsample for speed
 
 
def one_hot(y: np.ndarray, num_classes: int) -> np.ndarray:
    out = np.zeros((y.shape[0], num_classes), dtype=np.float32)
    out[np.arange(y.shape[0]), y] = 1.0
    return out
 
 
# ---------------- train loop ----------------
def train(mlp: MLP, Xtr, ytr, Xte, yte, *, epochs=15, batch=128, lr=0.1) -> None:
    ytr_oh = one_hot(ytr, 10)
    n = Xtr.shape[0]
    for ep in range(1, epochs + 1):
        idx = RNG.permutation(n)
        for start in range(0, n, batch):
            b_idx = idx[start:start + batch]
            probs, cache = mlp.forward(Xtr[b_idx])
            _, dW, db = mlp.loss_and_grads(probs, ytr_oh[b_idx], cache)
            mlp.step(dW, db, lr)
        probs_te, _ = mlp.forward(Xte)
        acc = float((probs_te.argmax(axis=1) == yte).mean())
        print(f"  epoch {ep:>2}: test accuracy = {acc:.4f}")
 
 
if __name__ == "__main__":
    Xtr, ytr, Xte, yte = load_mnist()
    print(f"Train: {Xtr.shape}  Test: {Xte.shape}")
 
    mlp = MLP(layer_sizes=[784, 128, 64, 10])
    total = sum(w.size for w in mlp.W) + sum(b.size for b in mlp.b)
    print(f"Model parameters: {total:,}")
 
    train(mlp, Xtr, ytr, Xte, yte, epochs=15, batch=128, lr=0.1)