Same Weights Everywhere: From Dense Layers to Convolutions

# Chapter 15: from dense layers to convolutions. Same weights at every position, so a shifted digit is still a digit.
# Uses tensor.py (Chapter 6), optim.py (Chapter 10's Adam), the mnist/ folder (Chapter 13)
# and digits_model.npz, the dense network saved in Chapter 14.
import gc
import gzip
import os
import struct
import time
import numpy as np
from tensor import Tensor
from optim import Adam

# === 1. Three new operations, each with its own backward rule =====================
def reshape(x, shape):
    out = Tensor(x.data.reshape(shape), (x,), "reshape")

    def _backward():                                   # same numbers, so the gradient just changes shape back
        x.grad += out.grad.reshape(x.shape)
    out._backward = _backward
    return out


def unfold(x, k):
    """(N, H, W, C) → (N, H−k+1, W−k+1, k·k·C): every k × k patch laid out as one row (often called im2col).
    Given a plain array, such as the input images, it skips the backward pass: data needs no gradient."""
    data = x.data if isinstance(x, Tensor) else x
    N, H, W, C = data.shape
    Ho, Wo = H - k + 1, W - k + 1
    windows = np.lib.stride_tricks.sliding_window_view(data, (k, k), axis=(1, 2))   # (N, Ho, Wo, C, k, k)
    cols = windows.transpose(0, 1, 2, 4, 5, 3).reshape(N, Ho, Wo, k * k * C)
    if not isinstance(x, Tensor):
        return Tensor(cols)
    out = Tensor(cols, (x,), "unfold")

    def _backward():                                   # each pixel appears in up to k·k patches: add them all up
        g = out.grad.reshape(N, Ho, Wo, k, k, C)
        for i in range(k):
            for j in range(k):
                x.grad[:, i:i + Ho, j:j + Wo, :] += g[:, :, :, i, j, :]
    out._backward = _backward
    return out


def maxpool2(x):
    """(N, H, W, C) → (N, H/2, W/2, C): keep the largest value in each 2 × 2 block."""
    N, H, W, C = x.shape
    blocks = x.data.reshape(N, H // 2, 2, W // 2, 2, C)
    biggest = blocks.max(axis=(2, 4))
    out = Tensor(biggest, (x,), "maxpool")

    def _backward():                                   # the gradient goes only to the winner of each block
        winner = blocks == biggest[:, :, None, :, None, :]
        winner = winner / winner.sum(axis=(2, 4), keepdims=True)
        x.grad += (winner * out.grad[:, :, None, :, None, :]).reshape(x.shape)
    out._backward = _backward
    return out


def conv(x, K, kb, k):
    """Convolution as a matrix product: every patch times every filter. K is (filters × k·k·channels)."""
    patches = unfold(x, k)
    N, Ho, Wo, P = patches.shape
    return reshape(reshape(patches, (N * Ho * Wo, P)) @ K.T + kb, (N, Ho, Wo, K.shape[0]))

# Gradient check on a tiny example (Chapter 4's habit, applied to the new operations)
r = np.random.default_rng(0)
x = Tensor(r.normal(size=(2, 8, 8, 3)))
Kc, kc = Tensor(r.normal(size=(4, 27))), Tensor(r.normal(size=4))
probe = r.normal(size=(2, 3, 3, 4))
f = lambda: (maxpool2(conv(x, Kc, kc, 3).relu()) * Tensor(probe)).sum()
f().backward()
worst = 0.0
for t, idx in [(x, (0, 2, 3, 1)), (x, (1, 5, 5, 0)), (Kc, (2, 10)), (Kc, (3, 26)), (kc, (1,))]:
    old = t.data[idx]
    t.data[idx] = old + 1e-6
    up = float(f().data)
    t.data[idx] = old - 1e-6
    down = float(f().data)
    t.data[idx] = old
    numeric = (up - down) / 2e-6
    worst = max(worst, abs(numeric - t.grad[idx]) / (abs(numeric) + abs(t.grad[idx])))
print(f"gradient check of unfold, conv, ReLU and max-pool: worst relative error {worst:.1e}")

# === 2. Why sharing weights matters ================================================
print(f"\none 5×5 filter: 25 weights + 1 bias, reused at all 24×24 = 576 positions")
print(f"8 filters: {8 * 26} knobs;  a dense layer making the same 24×24×8 outputs: {784 * 24 * 24 * 8:,} weights")

# Shift the input, and the feature map shifts with it (away from the edges)
img = r.random((1, 28, 28, 1))
moved = np.roll(img, 2, axis=2)
moved[:, :, :2] = 0
K1 = Tensor(r.normal(size=(1, 25)))
a = conv(img, K1, Tensor(np.zeros(1)), 5).data
b = conv(moved, K1, Tensor(np.zeros(1)), 5).data
print("conv(shifted image) equals shifted conv(image), away from the left edge:",
      np.allclose(b[:, :, 2:], a[:, :, :-2]))

# === 3. A small convolutional network for the digits =================================
def read_idx(name):
    with gzip.open(os.path.join("mnist", name), "rb") as fh:
        data = fh.read()
    ndim = data[3]
    shape = struct.unpack(">" + "I" * ndim, data[4:4 + 4 * ndim])
    return np.frombuffer(data, dtype=np.uint8, offset=4 + 4 * ndim).reshape(shape)

images, labels = read_idx("train-images-idx3-ubyte.gz"), read_idx("train-labels-idx1-ubyte.gz")
test_images, test_labels = read_idx("t10k-images-idx3-ubyte.gz"), read_idx("t10k-labels-idx1-ubyte.gz")
X_train, y_train = images[:50000, :, :, None] / 255.0, labels[:50000]     # keep the 28×28 shape, 1 channel
X_val, y_val = images[50000:, :, :, None] / 255.0, labels[50000:]

rng = np.random.default_rng(0)
K = Tensor(rng.normal(0, np.sqrt(2 / 25), (8, 25)))          # 8 filters of 5×5 (He initialization, Chapter 12)
kb = Tensor(np.zeros(8))
W = Tensor(rng.normal(0, np.sqrt(2 / 1152), (10, 1152)))      # 12×12×8 = 1,152 pooled features → 10 scores
b = Tensor(np.zeros(10))
params = [K, kb, W, b]
print(f"\nCNN: conv 8×5×5 → ReLU → 2×2 max-pool → dense 10;  {sum(p.data.size for p in params):,} knobs")


def scores(X):
    h = maxpool2(conv(X, K, kb, 5).relu())                    # (N, 12, 12, 8)
    return reshape(h, (len(X), 1152)) @ W.T + b


def softmax_cross_entropy(Z, y):
    E = (Z + (-Z.data.max(axis=1, keepdims=True))).exp()
    P = E / (E @ Tensor(np.ones((10, 1))))
    return -(Tensor(np.eye(10)[y]) * P.log()).sum() / Z.shape[0]


def accuracy(predict, X, y, batch=500):
    """In batches, collecting each batch's graph: the patches alone for 10,000 images would be about 1 GB."""
    right = 0
    for s in range(0, len(X), batch):
        right += int(np.sum(predict(X[s:s + batch]).argmax(axis=1) == y[s:s + batch]))
        gc.collect()
    return right / len(X)

opt = Adam(params, 0.001)
shuffle = np.random.default_rng(0)
for epoch in range(1, 6):
    t0 = time.time()
    order = shuffle.permutation(len(X_train))
    for start in range(0, len(X_train), 128):
        rows = order[start:start + 128]
        loss = softmax_cross_entropy(scores(X_train[rows]), y_train[rows])
        for p in params:
            p.grad = np.zeros_like(p.data)
        loss.backward()
        opt.step()
        del loss
        gc.collect()          # each step's graph holds reference cycles; free its big arrays now, not eventually
    print(f"  epoch {epoch}: validation {accuracy(lambda X: scores(X).data, X_val, y_val):.2%}  ({time.time() - t0:.0f} s)")
cnn = lambda X: scores(X).data
print(f"test accuracy: {accuracy(cnn, test_images[:, :, :, None] / 255.0, test_labels):.2%}")

# === 4. Chapter 13's shift test: the dense network from Chapter 14 against the CNN ====
with np.load("digits_model.npz") as saved:
    dense = {k: saved[k] for k in saved.files}
mlp = lambda X: np.maximum(0, X.reshape(len(X), 784) @ dense["W1"].T + dense["b1"]) @ dense["W2"].T + dense["b2"]


def shift_right(raw, pixels):
    out = np.roll(raw, pixels, axis=2)
    out[:, :, :pixels] = 0                                    # blank the columns that wrapped around
    return out

print("\nshift right by   dense network (101,770 knobs)   CNN (11,738 knobs)")
for pixels in [0, 1, 2, 3, 4]:
    X_shift = shift_right(test_images, pixels) / 255.0
    print(f"  {pixels} pixels        {accuracy(mlp, X_shift, test_labels):10.2%}"
          f"                   {accuracy(cnn, X_shift[:, :, :, None], test_labels):8.2%}")