#!/usr/bin/env python3 """ Tiny MNIST CNN in pure numpy — trains, evaluates, and exports weights as JSON for an in-browser digit recognizer. Architecture (kept deliberately small so the JS forward pass is a faithful, readable mirror and the weight file stays light): input 28x28x1 (x = pixel/255, then (x - 0.1307) / 0.3081 — std MNIST norm) conv1 3x3, 1->8, pad=1, stride=1 -> 28x28x8 -> ReLU pool1 2x2 max, stride=2 -> 14x14x8 conv2 3x3, 8->16, pad=1, stride=1 -> 14x14x16 -> ReLU pool2 2x2 max, stride=2 -> 7x7x16 flatten -> 784 dense 784 -> 10 -> softmax Trained with Adam + cross-entropy. ~9k params. Exports weights rounded to 5 sig figs as JSON (a few hundred KB), plus the exact preprocessing constants. """ import gzip, json, os, struct, sys, time, urllib.request import numpy as np np.random.seed(7) HERE = os.path.dirname(os.path.abspath(__file__)) CACHE = os.path.join(HERE, "mnist") os.makedirs(CACHE, exist_ok=True) BASE = "https://storage.googleapis.com/cvdf-datasets/mnist/" FILES = { "trX": "train-images-idx3-ubyte.gz", "trY": "train-labels-idx1-ubyte.gz", "teX": "t10k-images-idx3-ubyte.gz", "teY": "t10k-labels-idx1-ubyte.gz", } def fetch(name): path = os.path.join(CACHE, FILES[name]) if not os.path.exists(path): print(f"downloading {FILES[name]} ...", flush=True) urllib.request.urlretrieve(BASE + FILES[name], path) with gzip.open(path, "rb") as f: return f.read() def load_images(name): raw = fetch(name) magic, n, rows, cols = struct.unpack(">IIII", raw[:16]) assert magic == 2051 return np.frombuffer(raw, np.uint8, offset=16).reshape(n, rows, cols).astype(np.float32) def load_labels(name): raw = fetch(name) magic, n = struct.unpack(">II", raw[:8]) assert magic == 2049 return np.frombuffer(raw, np.uint8, offset=8).astype(np.int64) MEAN, STD = 0.1307, 0.3081 def norm(x): # x: (N,28,28) in [0,255] return ((x / 255.0) - MEAN) / STD # ----------------------------- im2col helpers ------------------------------ def im2col(x, kh, kw, pad, stride): # x: (N, C, H, W) -> cols: (N, C*kh*kw, OH*OW) N, C, H, W = x.shape OH = (H + 2 * pad - kh) // stride + 1 OW = (W + 2 * pad - kw) // stride + 1 xp = np.pad(x, ((0, 0), (0, 0), (pad, pad), (pad, pad))) cols = np.empty((N, C, kh, kw, OH, OW), dtype=x.dtype) for i in range(kh): ie = i + stride * OH for j in range(kw): je = j + stride * OW cols[:, :, i, j, :, :] = xp[:, :, i:ie:stride, j:je:stride] return cols.reshape(N, C * kh * kw, OH * OW), OH, OW def col2im(cols, x_shape, kh, kw, pad, stride, OH, OW): N, C, H, W = x_shape xp = np.zeros((N, C, H + 2 * pad, W + 2 * pad), dtype=cols.dtype) cols = cols.reshape(N, C, kh, kw, OH, OW) for i in range(kh): ie = i + stride * OH for j in range(kw): je = j + stride * OW xp[:, :, i:ie:stride, j:je:stride] += cols[:, :, i, j, :, :] return xp[:, :, pad:pad + H, pad:pad + W] if pad else xp # ------------------------------- layers ------------------------------------ class Conv: def __init__(self, cin, cout, k=3, pad=1, stride=1): self.k, self.pad, self.stride = k, pad, stride self.cin, self.cout = cin, cout fan_in = cin * k * k self.W = (np.random.randn(cout, cin * k * k) * np.sqrt(2.0 / fan_in)).astype(np.float32) self.b = np.zeros(cout, np.float32) self._init_adam() def _init_adam(self): self.mW = np.zeros_like(self.W); self.vW = np.zeros_like(self.W) self.mb = np.zeros_like(self.b); self.vb = np.zeros_like(self.b) def forward(self, x): self.x_shape = x.shape cols, OH, OW = im2col(x, self.k, self.k, self.pad, self.stride) self.cols = cols; self.OH, self.OW = OH, OW N = x.shape[0] out = np.einsum("oc,ncl->nol", self.W, cols) + self.b[None, :, None] return out.reshape(N, self.cout, OH, OW) def backward(self, dout): N = dout.shape[0] dout_r = dout.reshape(N, self.cout, self.OH * self.OW) self.db = dout_r.sum(axis=(0, 2)) self.dW = np.einsum("nol,ncl->oc", dout_r, self.cols) dcols = np.einsum("oc,nol->ncl", self.W, dout_r) return col2im(dcols, self.x_shape, self.k, self.k, self.pad, self.stride, self.OH, self.OW) class ReLU: def forward(self, x): self.mask = x > 0 return x * self.mask def backward(self, d): return d * self.mask class MaxPool: def __init__(self, k=2, stride=2): self.k, self.stride = k, stride def forward(self, x): N, C, H, W = x.shape self.x_shape = x.shape k, s = self.k, self.stride OH, OW = (H - k) // s + 1, (W - k) // s + 1 self.OH, self.OW = OH, OW cols = np.empty((N, C, k, k, OH, OW), dtype=x.dtype) for i in range(k): for j in range(k): cols[:, :, i, j, :, :] = x[:, :, i:i + s * OH:s, j:j + s * OW:s] cols = cols.reshape(N, C, k * k, OH, OW) self.argmax = cols.argmax(axis=2) return cols.max(axis=2) def backward(self, dout): N, C, H, W = self.x_shape k, s = self.k, self.stride # scatter each pooled gradient back to the arg-max position in its window dx = np.zeros(self.x_shape, dtype=dout.dtype) for i in range(k): for j in range(k): mask = (self.argmax == (i * k + j)).astype(dout.dtype) dx[:, :, i:i + s * self.OH:s, j:j + s * self.OW:s] += mask * dout return dx class Dense: def __init__(self, nin, nout): self.W = (np.random.randn(nout, nin) * np.sqrt(2.0 / nin)).astype(np.float32) self.b = np.zeros(nout, np.float32) self.mW = np.zeros_like(self.W); self.vW = np.zeros_like(self.W) self.mb = np.zeros_like(self.b); self.vb = np.zeros_like(self.b) def forward(self, x): self.x = x return x @ self.W.T + self.b def backward(self, dout): self.dW = dout.T @ self.x self.db = dout.sum(axis=0) return dout @ self.W def softmax_ce(logits, y): z = logits - logits.max(axis=1, keepdims=True) ez = np.exp(z) p = ez / ez.sum(axis=1, keepdims=True) N = logits.shape[0] loss = -np.log(p[np.arange(N), y] + 1e-12).mean() d = p.copy() d[np.arange(N), y] -= 1 return loss, d / N, p # ------------------------------- Adam step --------------------------------- def adam(layer, lr, t, b1=0.9, b2=0.999, eps=1e-8): for P, dP, m, v in (("W", "dW", "mW", "vW"), ("b", "db", "mb", "vb")): p = getattr(layer, P); g = getattr(layer, dP) mm = getattr(layer, m); vv = getattr(layer, v) mm[:] = b1 * mm + (1 - b1) * g vv[:] = b2 * vv + (1 - b2) * (g * g) mh = mm / (1 - b1 ** t); vh = vv / (1 - b2 ** t) p -= lr * mh / (np.sqrt(vh) + eps) # --------------------------------- model ----------------------------------- conv1, relu1, pool1 = Conv(1, 8), ReLU(), MaxPool() conv2, relu2, pool2 = Conv(8, 16), ReLU(), MaxPool() dense = Dense(7 * 7 * 16, 10) def forward(x): a = pool1.forward(relu1.forward(conv1.forward(x))) a = pool2.forward(relu2.forward(conv2.forward(a))) flat = a.reshape(a.shape[0], -1) return dense.forward(flat), a.shape def backward(dlogits, conv_shape): d = dense.backward(dlogits).reshape(conv_shape) d = conv2.backward(relu2.backward(pool2.backward(d))) d = conv1.backward(relu1.backward(pool1.backward(d))) # --------------------------------- train ----------------------------------- def main(): t0 = time.time() Xtr = norm(load_images("trX"))[:, None, :, :] # (60000,1,28,28) Ytr = load_labels("trY") Xte = norm(load_images("teX"))[:, None, :, :] Yte = load_labels("teY") print(f"data loaded {Xtr.shape} in {time.time()-t0:.1f}s", flush=True) N = Xtr.shape[0] bs = 64 lr = 1e-3 EPOCHS = int(os.environ.get("EPOCHS", "2")) TIME_BUDGET = float(os.environ.get("TIME_BUDGET", "1500")) # seconds t = 0 best_acc = 0.0 for ep in range(EPOCHS): idx = np.random.permutation(N) run_loss = 0.0 for b in range(0, N, bs): sel = idx[b:b + bs] xb, yb = Xtr[sel], Ytr[sel] logits, cshape = forward(xb) loss, dlogits, _ = softmax_ce(logits, yb) backward(dlogits, cshape) t += 1 for layer in (conv1, conv2, dense): adam(layer, lr, t) run_loss += loss if (b // bs) % 100 == 0: print(f"ep{ep} it{b//bs} loss {loss:.3f} " f"({time.time()-t0:.0f}s)", flush=True) if time.time() - t0 > TIME_BUDGET: print("time budget hit, stopping", flush=True) break # eval on a test subset for speed acc = evaluate(Xte, Yte) print(f"== epoch {ep}: test acc {acc*100:.2f}% " f"avg loss {run_loss/(b//bs+1):.3f} ({time.time()-t0:.0f}s)", flush=True) best_acc = max(best_acc, acc) if time.time() - t0 > TIME_BUDGET: break final = evaluate(Xte, Yte, full=True) print(f"FINAL full-test accuracy: {final*100:.2f}%", flush=True) export(final) def evaluate(X, Y, full=False): n = X.shape[0] if full else 2000 correct = 0 for b in range(0, n, 500): xb = X[b:b + 500] logits, _ = forward(xb) correct += (logits.argmax(1) == Y[b:b + 500]).sum() return correct / n def r5(a): # round to 5 sig figs, return nested python lists return np.round(a.astype(np.float64), 5).tolist() def export(acc): out = { "arch": "conv3x3(1->8,pad1)-relu-maxpool2 | conv3x3(8->16,pad1)-relu-maxpool2 | dense(784->10)-softmax", "norm": {"div": 255.0, "mean": MEAN, "std": STD}, "test_acc": round(float(acc), 4), # conv weights stored as (cout, cin, 3, 3) "conv1": {"w": r5(conv1.W.reshape(8, 1, 3, 3)), "b": r5(conv1.b)}, "conv2": {"w": r5(conv2.W.reshape(16, 8, 3, 3)), "b": r5(conv2.b)}, # dense stored as (10, 784); flatten order is channel-major: c*7*7 + y*7 + x "dense": {"w": r5(dense.W), "b": r5(dense.b)}, } path = os.path.join(HERE, "cnn_weights.json") with open(path, "w") as f: json.dump(out, f, separators=(",", ":")) kb = os.path.getsize(path) / 1024 print(f"wrote {path} ({kb:.0f} KB) test_acc={acc*100:.2f}%", flush=True) if __name__ == "__main__": main()