Real Data: Reading Handwritten Digits

# Chapter 13: real data. Handwritten digits (MNIST), from raw files to a network that reads them.
# Uses tensor.py (Chapter 6) and optim.py (Chapter 10's Adam, as in Chapter 11).
# MNIST is by Yann LeCun, Corinna Cortes and Christopher Burges, licensed CC BY-SA 3.0.
import gzip
import hashlib
import os
import struct
import urllib.request
import numpy as np
from tensor import Tensor
from optim import Adam

# === 1. Get the files, and check they are the right ones ==========================
MIRROR = "https://storage.googleapis.com/cvdf-datasets/mnist/"
FILES = {"train-images-idx3-ubyte.gz": "440fcabf73cc546fa21475e81ea370265605f56be210a4024d2ca8f203523609",
         "train-labels-idx1-ubyte.gz": "3552534a0a558bbed6aed32b30c495cca23d567ec52cac8be1a0730e8010255c",
         "t10k-images-idx3-ubyte.gz": "8d422c7b0a1c1c79245a5bcf07fe86e33eeafee792b84584aec276f5a2dbc4e6",
         "t10k-labels-idx1-ubyte.gz": "f7ae60f92e00ec6debd23a6088c31dbd2371eca3ffa0defaefb259924204aec6"}
os.makedirs("mnist", exist_ok=True)
for name, sha in FILES.items():
    path = os.path.join("mnist", name)
    if not os.path.exists(path):
        urllib.request.urlretrieve(MIRROR + name, path)
    with open(path, "rb") as f:
        assert hashlib.sha256(f.read()).hexdigest() == sha, f"{name} is damaged or not the file we expect"


def read_idx(name):
    """IDX format: two zero bytes, a type code, the number of dimensions, each size (big-endian), then raw bytes."""
    with gzip.open(os.path.join("mnist", name), "rb") as f:
        data = f.read()
    _zero, _type, ndim = struct.unpack(">HBB", data[:4])
    shape = struct.unpack(">" + "I" * ndim, data[4:4 + 4 * ndim])
    return np.frombuffer(data, dtype=np.uint8, offset=4 + 4 * ndim).reshape(shape)

train_images, train_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")
print("train:", train_images.shape, " test:", test_images.shape, " pixel values", train_images.min(), "to", train_images.max())
print("training examples per digit:", np.bincount(train_labels).tolist())

# === 2. Look at one before doing anything else =====================================
print(f"\nthe first training image, labelled {train_labels[0]}:")
for row in train_images[0][4:24]:
    print("   " + "".join(" .:#"[v // 64] for v in row[3:25]))

# === 3. Flatten, scale, split =======================================================
X_all = train_images.reshape(-1, 784) / 255.0      # each image becomes 784 numbers between 0 and 1
X_test = test_images.reshape(-1, 784) / 255.0
X_train, y_train = X_all[:50000], train_labels[:50000]
X_val, y_val = X_all[50000:], train_labels[50000:]  # the last 10,000 training images, for choosing (Chapter 11)
print(f"\nmean pixel {X_train.mean():.4f};  50,000 training, 10,000 validation, 10,000 test")
print(f"always guessing the commonest digit: {np.bincount(y_val).max() / len(y_val):.1%} of validation images;"
      f" loss of a uniform guess: ln 10 = {np.log(10):.4f}")


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


def evaluate(model, X, labels):
    loss, P = softmax_cross_entropy(model(X), labels)
    return float(loss.data), float(np.mean(P.data.argmax(axis=1) == labels)), P.data.argmax(axis=1)


def train(params, model, epochs, lr=0.001, batch=128, seed=0):
    """Mini-batch Adam (Chapters 7 and 10). Keeps the weights from the epoch with the lowest validation loss."""
    opt = Adam(params, lr)
    order = np.random.default_rng(seed)
    best = (np.inf, None, 0)
    for epoch in range(1, epochs + 1):
        shuffled = order.permutation(len(X_train))
        for start in range(0, len(X_train), batch):
            rows = shuffled[start:start + batch]
            loss, _ = softmax_cross_entropy(model(X_train[rows]), y_train[rows])
            for p in params:
                p.grad = np.zeros_like(p.data)
            loss.backward()
            opt.step()
        tl, ta, _ = evaluate(model, X_train, y_train)
        vl, va, _ = evaluate(model, X_val, y_val)
        print(f"  epoch {epoch:2}: train {tl:.4f} ({ta:.2%})   validation {vl:.4f} ({va:.2%})")
        if vl < best[0]:
            best = (vl, [p.data.copy() for p in params], epoch)
    for p, saved in zip(params, best[1]):
        p.data = saved
    return best[2]

# === 4. Softmax regression: one score per digit, straight from the pixels ===========
W = Tensor(np.zeros((10, 784)))
b = Tensor(np.zeros(10))
linear = lambda X: Tensor(X) @ W.T + b
print("\nsoftmax regression (7,850 knobs):")
kept = train([W, b], linear, epochs=5)
loss, acc, _ = evaluate(linear, X_test, test_labels)
print(f"kept epoch {kept}; test: loss {loss:.4f}, accuracy {acc:.2%}")

# === 5. One hidden layer of 128 ReLU units, He initialization (Chapter 12) ===========
rng = np.random.default_rng(0)
W1, b1 = Tensor(rng.normal(0, np.sqrt(2 / 784), (128, 784))), Tensor(np.zeros(128))
W2, b2 = Tensor(rng.normal(0, np.sqrt(2 / 128), (10, 128))), Tensor(np.zeros(10))
mlp = lambda X: (Tensor(X) @ W1.T + b1).relu() @ W2.T + b2
print("\n784-128-10 network (101,770 knobs):")
kept = train([W1, b1, W2, b2], mlp, epochs=10)
loss, acc, predicted = evaluate(mlp, X_test, test_labels)
print(f"kept epoch {kept}; test: loss {loss:.4f}, accuracy {acc:.2%}, {np.sum(predicted != test_labels)} mistakes in 10,000")

# === 6. Which digits does it mix up? ===================================================
confusion = np.zeros((10, 10), dtype=int)
np.add.at(confusion, (test_labels, predicted), 1)
np.fill_diagonal(confusion, 0)
pairs = sorted(((confusion[t, p], t, p) for t in range(10) for p in range(10) if t != p), reverse=True)[:6]
print("most common mistakes (true → predicted, count):", ", ".join(f"{t}→{p}: {n}" for n, t, p in pairs))
print("accuracy per digit:", [f"{d}: {1 - confusion[d].sum() / np.sum(test_labels == d):.1%}" for d in range(10)])