Ch 1 — Neural Networks from Scratch¶

"""
Two-layer MLP implemented from scratch in NumPy.
Architecture: Linear(d_in, d_hidden) -> ReLU -> Linear(d_hidden, d_out) -> Softmax
Loss: Cross-entropy
Optimiser: Mini-batch SGD
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray


# ---------------------------------------------------------------------------
# Activation functions
# ---------------------------------------------------------------------------

def relu(z: NDArray[np.float64]) -> NDArray[np.float64]:
    """Rectified linear unit: max(0, z)."""
    return np.maximum(0.0, z)


def relu_grad(z: NDArray[np.float64]) -> NDArray[np.float64]:
    """Element-wise gradient of ReLU."""
    return (z > 0).astype(np.float64)


def softmax(z: NDArray[np.float64]) -> NDArray[np.float64]:
    """Numerically stable row-wise softmax.

    Args:
        z: Pre-softmax logits, shape (N, C).

    Returns:
        Probability matrix, shape (N, C). Each row sums to 1.
    """
    shifted = z - z.max(axis=1, keepdims=True)
    exp_z = np.exp(shifted)
    return exp_z / exp_z.sum(axis=1, keepdims=True)


# ---------------------------------------------------------------------------
# Loss
# ---------------------------------------------------------------------------

def cross_entropy_loss(
    y_hat: NDArray[np.float64],
    y: NDArray[np.int64],
) -> float:
    """Categorical cross-entropy loss.

    Args:
        y_hat: Predicted probabilities, shape (N, C).
        y:     Ground-truth class indices, shape (N,).

    Returns:
        Scalar mean loss.
    """
    n = y_hat.shape[0]
    log_probs = np.log(y_hat[np.arange(n), y] + 1e-12)
    return -float(log_probs.mean())


# ---------------------------------------------------------------------------
# MLP
# ---------------------------------------------------------------------------

class MLP:
    """Two-layer MLP with ReLU hidden activation and softmax output.

    Args:
        d_in:     Input feature dimension.
        d_hidden: Number of hidden units.
        d_out:    Number of output classes.
        lr:       Learning rate for SGD.
        seed:     Random seed for reproducibility.
    """

    def __init__(
        self,
        d_in: int,
        d_hidden: int,
        d_out: int,
        lr: float = 0.01,
        seed: int = 42,
    ) -> None:
        rng = np.random.default_rng(seed)

        # He initialisation for ReLU-activated hidden layer
        self.W1: NDArray[np.float64] = rng.normal(
            0.0, np.sqrt(2.0 / d_in), size=(d_hidden, d_in)
        )
        self.b1: NDArray[np.float64] = np.zeros(d_hidden)

        # Xavier initialisation for output layer (no ReLU)
        self.W2: NDArray[np.float64] = rng.normal(
            0.0, np.sqrt(1.0 / d_hidden), size=(d_out, d_hidden)
        )
        self.b2: NDArray[np.float64] = np.zeros(d_out)

        self.lr = lr

        # Cache for backprop
        self._x: NDArray[np.float64] | None = None
        self._z1: NDArray[np.float64] | None = None
        self._h1: NDArray[np.float64] | None = None
        self._z2: NDArray[np.float64] | None = None
        self._y_hat: NDArray[np.float64] | None = None

    # ------------------------------------------------------------------
    # Forward pass
    # ------------------------------------------------------------------

    def forward(self, x: NDArray[np.float64]) -> NDArray[np.float64]:
        """Run a forward pass.

        Args:
            x: Input batch, shape (N, d_in).

        Returns:
            Softmax probabilities, shape (N, d_out).
        """
        self._x = x                                  # (N, d_in)

        self._z1 = x @ self.W1.T + self.b1           # (N, d_hidden)
        self._h1 = relu(self._z1)                    # (N, d_hidden)

        self._z2 = self._h1 @ self.W2.T + self.b2   # (N, d_out)
        self._y_hat = softmax(self._z2)              # (N, d_out)

        return self._y_hat

    # ------------------------------------------------------------------
    # Backward pass
    # ------------------------------------------------------------------

    def backward(self, y: NDArray[np.int64]) -> None:
        """Compute gradients and update parameters.

        Args:
            y: Ground-truth class indices, shape (N,).
        """
        n = self._y_hat.shape[0]  # type: ignore[union-attr]

        # Output layer delta: softmax + cross-entropy combined gradient
        delta2 = self._y_hat.copy()                  # (N, d_out)
        delta2[np.arange(n), y] -= 1.0
        delta2 /= n

        # Gradients for W2, b2
        dW2 = delta2.T @ self._h1                    # (d_out, d_hidden)
        db2 = delta2.sum(axis=0)                     # (d_out,)

        # Propagate through hidden layer
        delta1 = delta2 @ self.W2                    # (N, d_hidden)
        delta1 *= relu_grad(self._z1)               # element-wise mask

        # Gradients for W1, b1
        dW1 = delta1.T @ self._x                    # (d_hidden, d_in)
        db1 = delta1.sum(axis=0)                    # (d_hidden,)

        # SGD update
        self.W1 -= self.lr * dW1
        self.b1 -= self.lr * db1
        self.W2 -= self.lr * dW2
        self.b2 -= self.lr * db2

    # ------------------------------------------------------------------
    # Training loop
    # ------------------------------------------------------------------

    def fit(
        self,
        x_train: NDArray[np.float64],
        y_train: NDArray[np.int64],
        epochs: int = 100,
        batch_size: int = 32,
        verbose: bool = True,
    ) -> list[float]:
        """Train the MLP using mini-batch SGD.

        Args:
            x_train:    Training inputs, shape (N, d_in).
            y_train:    Training labels, shape (N,).
            epochs:     Number of full passes over the data.
            batch_size: Mini-batch size.
            verbose:    Print loss every 10 epochs.

        Returns:
            List of per-epoch mean losses.
        """
        n = x_train.shape[0]
        history: list[float] = []
        rng = np.random.default_rng(0)

        for epoch in range(epochs):
            # Shuffle data
            perm = rng.permutation(n)
            x_shuf, y_shuf = x_train[perm], y_train[perm]

            epoch_losses: list[float] = []
            for start in range(0, n, batch_size):
                xb = x_shuf[start : start + batch_size]
                yb = y_shuf[start : start + batch_size]

                y_hat = self.forward(xb)
                loss = cross_entropy_loss(y_hat, yb)
                self.backward(yb)

                epoch_losses.append(loss)

            mean_loss = float(np.mean(epoch_losses))
            history.append(mean_loss)

            if verbose and (epoch + 1) % 10 == 0:
                print(f"Epoch {epoch + 1:4d}/{epochs}  loss={mean_loss:.4f}")

        return history

    def predict(self, x: NDArray[np.float64]) -> NDArray[np.int64]:
        """Return predicted class indices."""
        probs = self.forward(x)
        return np.argmax(probs, axis=1).astype(np.int64)


# ---------------------------------------------------------------------------
# Demo
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    from sklearn.datasets import make_classification
    from sklearn.model_selection import train_test_split
    from sklearn.preprocessing import StandardScaler

    X, y = make_classification(
        n_samples=1000, n_features=10, n_classes=3,
        n_informative=6, random_state=0
    )
    X = StandardScaler().fit_transform(X)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

    model = MLP(d_in=10, d_hidden=64, d_out=3, lr=0.05)
    history = model.fit(X_train, y_train, epochs=50, batch_size=32)

    preds = model.predict(X_test)
    acc = (preds == y_test).mean()
    print(f"\nTest accuracy: {acc:.3f}")

"""
Two-layer MLP implemented in PyTorch using nn.Module.
Equivalent architecture to the NumPy version above.
"""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.utils.data import DataLoader, TensorDataset


class MLP(nn.Module):
    """Two-layer MLP with ReLU hidden activation.

    Args:
        d_in:     Input feature dimension.
        d_hidden: Number of hidden units.
        d_out:    Number of output classes.
        dropout:  Dropout probability applied after hidden layer.
    """

    def __init__(
        self,
        d_in: int,
        d_hidden: int,
        d_out: int,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in, d_hidden),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_hidden, d_out),
        )
        self._init_weights()

    def _init_weights(self) -> None:
        """Apply He and Xavier initialisation to linear layers."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
                nn.init.zeros_(module.bias)
        # Output layer: Xavier
        nn.init.xavier_uniform_(self.net[-1].weight)  # type: ignore[arg-type]

    def forward(self, x: Tensor) -> Tensor:
        """Forward pass returning logits (pre-softmax).

        Args:
            x: Input batch, shape (N, d_in).

        Returns:
            Logits, shape (N, d_out).
        """
        return self.net(x)


def train_one_epoch(
    model: MLP,
    loader: DataLoader,
    optimiser: torch.optim.Optimizer,
    device: torch.device,
) -> float:
    """Train for one epoch and return mean loss."""
    model.train()
    total_loss = 0.0
    for x_batch, y_batch in loader:
        x_batch, y_batch = x_batch.to(device), y_batch.to(device)

        optimiser.zero_grad()
        logits = model(x_batch)
        loss = F.cross_entropy(logits, y_batch)
        loss.backward()
        optimiser.step()

        total_loss += loss.item() * x_batch.size(0)

    return total_loss / len(loader.dataset)  # type: ignore[arg-type]


@torch.no_grad()
def evaluate(
    model: MLP,
    loader: DataLoader,
    device: torch.device,
) -> tuple[float, float]:
    """Return (mean_loss, accuracy) on a dataset."""
    model.eval()
    total_loss, correct = 0.0, 0
    for x_batch, y_batch in loader:
        x_batch, y_batch = x_batch.to(device), y_batch.to(device)
        logits = model(x_batch)
        total_loss += F.cross_entropy(logits, y_batch, reduction="sum").item()
        correct += (logits.argmax(dim=1) == y_batch).sum().item()

    n = len(loader.dataset)  # type: ignore[arg-type]
    return total_loss / n, correct / n


if __name__ == "__main__":
    import numpy as np
    from sklearn.datasets import make_classification
    from sklearn.model_selection import train_test_split
    from sklearn.preprocessing import StandardScaler

    # --- Data ---
    X, y = make_classification(
        n_samples=1000, n_features=10, n_classes=3,
        n_informative=6, random_state=0
    )
    X = StandardScaler().fit_transform(X).astype(np.float32)
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    train_ds = TensorDataset(torch.from_numpy(X_tr), torch.tensor(y_tr))
    test_ds  = TensorDataset(torch.from_numpy(X_te), torch.tensor(y_te))
    train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
    test_loader  = DataLoader(test_ds,  batch_size=64)

    # --- Model ---
    model = MLP(d_in=10, d_hidden=64, d_out=3, dropout=0.1).to(device)
    optimiser = torch.optim.Adam(model.parameters(), lr=1e-3)

    # --- Training ---
    for epoch in range(50):
        train_loss = train_one_epoch(model, train_loader, optimiser, device)
        if (epoch + 1) % 10 == 0:
            _, acc = evaluate(model, test_loader, device)
            print(f"Epoch {epoch+1:3d}  train_loss={train_loss:.4f}  test_acc={acc:.3f}")