#!/usr/bin/env python3 """ nnet - train a small MLP on MNIST and quantise it for integer inference (offline). This runs off the machine (on the Linux host) and is the only part of the project that uses floating point. It writes: src/model.h - the quantised network as C89 arrays (int16 weights, int32 biases, one power-of-two shift), plus the exact integer inference in numpy so its accuracy can be checked before a line of C is written. data/nn_test.bin - the first N test images and labels, a compact binary the 604e reads at run time. Quantisation scheme (fully integer inference): * Inputs are raw pixels 0..255 (the 1/255 normalisation is folded into the first-layer weights). * Each layer's weights are symmetric int16 with a per-layer power-of-two quantum uw = 2^ceil(log2(max|W|/32767)). * Biases are folded in as int32: bq = round(b / (ux*uw)). * Accumulate in int64, then the hidden layer is requantised with a single arithmetic right shift (plus rounding) and ReLU; the output layer is just the argmax of the integer accumulators (scaling does not change argmax). DeepSeek (deepseek-v4.1-flash), October 2026. """ import gzip import os import struct import sys import numpy as np RNG = np.random.default_rng(20261015) DATA = "data" HID = 32 EPOCHS = 15 BATCH = 128 LR = 1e-3 N_TEST_ON_MACHINE = 1000 def load_images(path): with gzip.open(path, "rb") as f: magic, n, rows, cols = struct.unpack(">IIII", f.read(16)) assert magic == 2051, magic buf = f.read(n * rows * cols) return np.frombuffer(buf, dtype=np.uint8).reshape(n, rows * cols).copy() def load_labels(path): with gzip.open(path, "rb") as f: magic, n = struct.unpack(">II", f.read(8)) assert magic == 2049, magic buf = f.read(n) return np.frombuffer(buf, dtype=np.uint8).copy() def train(x, y): n, d = x.shape w1 = RNG.standard_normal((d, HID)) * np.sqrt(2.0 / d) b1 = np.zeros(HID) w2 = RNG.standard_normal((HID, 10)) * np.sqrt(2.0 / HID) b2 = np.zeros(10) params = [w1, b1, w2, b2] m = [np.zeros_like(p) for p in params] v = [np.zeros_like(p) for p in params] b1p, b2p = 0.9, 0.999 eps = 1e-8 for ep in range(EPOCHS): idx = RNG.permutation(n) correct = 0 for s in range(0, n, BATCH): bi = idx[s:s + BATCH] xb = x[bi] / 255.0 yb = y[bi] z1 = xb @ w1 + b1 h1 = np.maximum(z1, 0.0) z2 = h1 @ w2 + b2 z2 -= z2.max(axis=1, keepdims=True) p = np.exp(z2) p /= p.sum(axis=1, keepdims=True) correct += (p.argmax(axis=1) == yb).sum() t = np.zeros((len(bi), 10)) t[np.arange(len(bi)), yb] = 1.0 dz2 = (p - t) / len(bi) dw2 = h1.T @ dz2 db2 = dz2.sum(axis=0) dh1 = dz2 @ w2.T dz1 = dh1 * (z1 > 0) dw1 = xb.T @ dz1 db1 = dz1.sum(axis=0) grads = [dw1, db1, dw2, db2] lr = LR * (0.5 ** (ep / 6.0)) for i, (p_, g) in enumerate(zip(params, grads)): m[i] = b1p * m[i] + (1 - b1p) * g v[i] = b2p * v[i] + (1 - b2p) * g * g mhat = m[i] / (1 - b1p ** (ep * (n // BATCH) + s // BATCH + 1)) vhat = v[i] / (1 - b2p ** (ep * (n // BATCH) + s // BATCH + 1)) p_ -= lr * mhat / (np.sqrt(vhat) + eps) print("epoch %2d train acc %.4f" % (ep + 1, correct / n)) return w1, b1, w2, b2 def pow2_quant(w, levels=32767): """Symmetric quantisation to int16 with a power-of-two quantum.""" mx = np.max(np.abs(w)) if mx == 0: return np.zeros_like(w, dtype=np.int16), 1.0, 0 # smallest power of two uw so that max|w|/uw <= levels k = int(np.ceil(np.log2(mx / levels))) uw = 2.0 ** k q = np.rint(w / uw).astype(np.int64) q = np.clip(q, -32768, 32767) return q.astype(np.int16), uw, k def quantise_and_simulate(w1, b1, w2, b2, x_train, x_test, y_test): """Fold the 1/255 into w1, quantise, then run the exact integer inference.""" w1f = w1 / 255.0 # inputs become raw pixels 0..255 qw1, uw1, _ = pow2_quant(w1f) qw2, uw2, _ = pow2_quant(w2) ux0 = 1.0 # input quantum: one pixel unit # calibrate the hidden post-ReLU range on a subset h_cal = np.maximum((x_train[:2000] / 255.0) @ w1 + b1, 0.0) hmax = float(h_cal.max()) uy0 = 2.0 ** int(np.ceil(np.log2(max(hmax, 1e-9) / 32767))) # integer biases, folded to the accumulator scale (ux*uw) bq0 = np.rint(b1 / (ux0 * uw1)).astype(np.int64) bq1 = np.rint(b2 / (uy0 * uw2)).astype(np.int64) # shift for the hidden requantisation: yq = round(acc * (ux0*uw1/uy0)) # ratio is a power of two: ratio = 2^m, so a left shift if m>=0 else right. ratio = (ux0 * uw1) / uy0 m = int(round(np.log2(ratio))) assert abs(ratio - 2.0 ** m) < 1e-12, ratio rs = -m if m < 0 else 0 def infer(xpix): # xpix: (n,784) uint8 -> predicted labels a = xpix.astype(np.int64) # (n,784) acc0 = a @ qw1.astype(np.int64) + bq0 # (n,32) pre = np.maximum(acc0, 0) # ReLU in the accumulator domain if m >= 0: h = pre << m else: h = (pre + (1 << (rs - 1))) >> rs h = np.clip(h, 0, 32767) # int16 backstop acc1 = h @ qw2.astype(np.int64) + bq1 # (n,10) return acc1.argmax(axis=1), acc1, h pred, _, _ = infer(x_test) acc = (pred == y_test).mean() # float reference on raw pixels (equivalent to the trained net) z1 = (x_test / 255.0) @ w1 + b1 h1 = np.maximum(z1, 0.0) z2 = h1 @ w2 + b2 facc = (z2.argmax(axis=1) == y_test).mean() print("float accuracy: %.4f" % facc) print("int accuracy: %.4f" % acc) print("quanta: uw1=2^%d uw2=2^%d uy0=2^%d shift=%d" % (int(round(np.log2(uw1))), int(round(np.log2(uw2))), int(round(np.log2(uy0))), m)) return qw1, bq0, qw2, bq1, m, acc, facc def write_model(qw1, bq0, qw2, bq1, shift, path): with open(path, "w") as f: f.write("/* generated by train.py - do not edit */\n") f.write("#ifndef NNET_MODEL_H\n#define NNET_MODEL_H\n\n") f.write("#define NN_IN %d\n#define NN_HID %d\n#define NN_OUT %d\n" % (qw1.shape[0], qw1.shape[1], qw2.shape[1])) f.write("#define NN_SHIFT_H %d\n\n" % shift) f.write("static const short nn_w0[NN_HID][NN_IN] = {\n") for j in range(qw1.shape[1]): f.write(" {" + ",".join(str(int(v)) for v in qw1[:, j]) + "},\n") f.write("};\n\n") f.write("static const long nn_b0[NN_HID] = {\n ") f.write(",".join(str(int(v)) for v in bq0) + "};\n\n") f.write("static const short nn_w1[NN_OUT][NN_HID] = {\n") for o in range(qw2.shape[1]): f.write(" {" + ",".join(str(int(v)) for v in qw2[:, o]) + "},\n") f.write("};\n\n") f.write("static const long nn_b1[NN_OUT] = {\n ") f.write(",".join(str(int(v)) for v in bq1) + "};\n\n") f.write("#endif /* NNET_MODEL_H */\n") def write_test_bin(x, y, path, n): with open(path, "wb") as f: f.write(struct.pack("