MNIST
The fundamental motivation here is to build the simplest possible, no frills, python standard libs only implementation of an MNIST classifier. By implementing everything from scratch and using very basic components, you can get a visceral appreciation for the utility of a whole range of things that you might otherwise take for granted: pytorch, GPUs, better activation functions, etc. I also believe that an ultra-simple implementation like this serves to demystify some ML concepts on a basic level. Sure, technical non-ML people might have an abstract understanding of backpropogation, weights, activations, etc, but seeing the structure laid out at an extremely basic level may prove useful. I thought about using numpy, and ultimately that implementation would have been fairly similar, however I really wanted to lean into the purist approach. Also by not using numpy, the utility of vector operations and the preponderance of linear algebra can be fully appreciated. mnist_torch.py is provided as juxtaposition to show what a comparable implementation looks like using pytorch & GPU acceleration.
Notes
The only conceptually annoying part of this process is keeping track of indexes and knowing how to calculate gradients during the backwards pass. otherwise the logic is trivial. To understand the structure of the weights, see below:
for hl_dims = [16, 16] # 2 hidden layers with 16 neurons each
Ws = [[w(0,0), w(1,0), ... w(14,0), w(15,0)],
[w(0,1), w(1,1), ... ],
...]
with w(x, y) where x = node_idx in last layer and y = node_idx in next layer
import json
import logging
import math
import os
import random
import struct
import time
import utils
DATA_CACHE_PATH = "mnist/.local/data/"
CHECKPOINT_PATH = "mnist/.local/checkpoints/"
TRAIN_IMG_PATH = f"{DATA_CACHE_PATH}train-images.idx3-ubyte"
TRAIN_LBL_PATH = f"{DATA_CACHE_PATH}train-labels.idx1-ubyte"
TEST_IMG_PATH = f"{DATA_CACHE_PATH}t10k-images.idx3-ubyte"
TEST_LBL_PATH = f"{DATA_CACHE_PATH}t10k-labels.idx1-ubyte"
NUM_TRAINING_EXAMPLES = 60_000
NUM_TESTING_EXAMPLES = 10_000
def main() -> None:
utils.download_dataset(
"https://www.kaggle.com/api/v1/datasets/download/hojjatk/mnist-dataset",
f"{DATA_CACHE_PATH}mnist.zip",
)
os.makedirs(DATA_CACHE_PATH, exist_ok=True)
os.system(f"unzip -q {DATA_CACHE_PATH}mnist.zip -d {DATA_CACHE_PATH}")
train(hidden_layer_dims=[16, 16], learning_rate=1e-2, iterations=50_000)
def train(
hidden_layer_dims: list[int],
learning_rate: float,
iterations: int,
) -> None:
start_ts = time.time()
checkpoint_fp = CHECKPOINT_PATH + "-".join(str(d) for d in hidden_layer_dims) + ".json"
sc = (
ScratchClassifier.load(checkpoint_fp)
if os.path.exists(checkpoint_fp)
else ScratchClassifier(hidden_layer_dims=hidden_layer_dims)
)
logging.info(f"training: {hidden_layer_dims=} | {iterations=} | {learning_rate=}")
losses = []
for iteration in range(iterations):
img, lbl = mnist_sample(random.randint(0, NUM_TRAINING_EXAMPLES - 1))
activations = sc.forward(img)
sc.backward(lbl, activations, learning_rate)
losses.append(-math.log(activations[-1][lbl] + 1e-12))
if iteration % 1_000 == 0 and iteration > 0:
logging.info(
f"{iteration=} - loss={sum(losses[-1000:]) / 1000:.4f} | elapsed time={time.time() - start_ts:.2f} secs"
)
logging.info(
f"{iteration=} - loss={sum(losses[-1000:]) / 1000:.4f} | elapsed time={time.time() - start_ts:.2f} secs - COMPLETE"
)
sc.save()
for iteration in range(10):
img, _ = mnist_sample(random.randint(0, NUM_TESTING_EXAMPLES - 1), train=False)
activations = sc.forward(img)
predicted_label = activations[-1].index(max(activations[-1]))
log_ascii_mnist(img)
logging.info(f"{predicted_label=}")
class ScratchClassifier:
def __init__(
self,
in_dim: int = 784,
hidden_layer_dims: list[int] = [16, 16],
out_dim: int = 10,
weights: list[list[float]] | None = None,
biases: list[list[float]] | None = None,
) -> None:
assert len(hidden_layer_dims) >= 1
self.in_dim, self.hl_dims, self.out_dim = in_dim, hidden_layer_dims, out_dim
self.Ws = weights or (
[xavier(self.in_dim, self.hl_dims[0])]
+ [xavier(self.hl_dims[i], self.hl_dims[i + 1]) for i in range(len(self.hl_dims) - 1)]
+ [xavier(self.hl_dims[-1], self.out_dim)]
)
self.Bs = biases or [[0.0] * d for d in self.hl_dims] + [[0.0] * self.out_dim]
def forward(self, img: list[int]) -> list[list[float]]:
"""returns list of activations at each layer"""
activations = [[px / 255.0 for px in img]]
for layer_idx in range(len(self.hl_dims)):
activations.append(sigmoid(self.forward_layer(activations[-1], layer_idx)))
# we want probabilities that sum to 1 in the final layer, so we use softmax instead of sigmoid
activations.append(softmax(self.forward_layer(activations[-1], layer_idx + 1)))
return activations
def forward_layer(
self,
prev_activations: list[float],
layer_idx: int,
) -> list[float]:
"""returns current layer activations pre-activation function"""
weighted_inputs = []
in_dim = len(prev_activations)
out_dim = len(self.Bs[layer_idx])
for node_idx in range(out_dim):
cur_weights = self.Ws[layer_idx][node_idx * in_dim : (node_idx + 1) * in_dim]
w_mul_a = [w * a for w, a in zip(cur_weights, prev_activations)]
weighted_inputs.append(sum(w_mul_a) + self.Bs[layer_idx][node_idx])
return weighted_inputs
def backward(self, label: int, activations: list[list[float]], learning_rate: float) -> None:
y_hat = activations[-1]
y = [1 if i == label else 0 for i in range(len(y_hat))]
deltas = [(a_i - y_i) for (a_i, y_i) in zip(y_hat, y)]
for layer_idx in reversed(range(len(self.Ws))):
deltas = self.backward_layer(layer_idx, deltas, activations[layer_idx], learning_rate)
def backward_layer(
self,
layer_idx: int,
deltas: list[float],
prev_activations: list[float],
learning_rate: float,
) -> list[float]:
"""returns list of deltas to pass backwards"""
inp_dim = len(prev_activations)
prev_deltas = [0.0] * inp_dim
for node_idx, delta in enumerate(deltas):
self.Bs[layer_idx][node_idx] -= learning_rate * delta
for prev_node_idx, prev_activation in enumerate(prev_activations):
w_idx = (node_idx * inp_dim) + prev_node_idx
prev_deltas[prev_node_idx] += self.Ws[layer_idx][w_idx] * delta
self.Ws[layer_idx][w_idx] -= learning_rate * delta * prev_activation
# don't pass deltas back on the first layer
return (
[]
if layer_idx == 0
else [d * sigmoid_deriv(a) for d, a in zip(prev_deltas, prev_activations)]
)
def save(self) -> None:
fp = CHECKPOINT_PATH + "-".join(str(d) for d in self.hl_dims) + ".json"
os.makedirs(CHECKPOINT_PATH, exist_ok=True)
with open(fp, "w") as f:
json.dump({"weights": self.Ws, "biases": self.Bs}, f)
@staticmethod
def load(fp: str | None = None) -> "ScratchClassifier":
fp = fp or CHECKPOINT_PATH + "16-16.json"
logging.info(f"loading checkpoint with {fp=}")
with open(fp, "r") as f:
checkpoint = json.load(f)
return ScratchClassifier(weights=checkpoint["weights"], biases=checkpoint["biases"])
def sigmoid(xs: list[float]) -> list[float]:
return [1 / (1 + math.exp(-x)) for x in xs]
def sigmoid_deriv(x: float) -> float:
return x * (1.0 - x)
def softmax(xs: list[float]) -> list[float]:
max_x = max(xs)
exps = [math.exp(x - max_x) for x in xs]
s = sum(exps)
return [e / s for e in exps]
def xavier(inp_dim: int, output_dim: int) -> list[float]:
limit = (6 / (inp_dim + output_dim)) ** 0.5
return [random.uniform(-limit, limit) for _ in range(inp_dim * output_dim)]
# HELPERS =====================================================================================
def mnist_sample(idx: int = 0, train: bool = True) -> tuple[list[int], int]:
img_path, lbl_path = (
(TRAIN_IMG_PATH, TRAIN_LBL_PATH) if train else (TEST_IMG_PATH, TEST_LBL_PATH)
)
with open(lbl_path, "rb") as f:
f.seek(8 + idx)
label = f.read(1)[0]
with open(img_path, "rb") as f:
_, _, rows, cols = struct.unpack(">IIII", f.read(16))
f.seek(idx * rows * cols, 1)
image = list(f.read(rows * cols))
return image, label
def log_ascii_mnist(img: list[int], w: int = 28) -> None:
chars = " ░▒▓█"
scale = len(chars) - 1
ascii_str = ""
for i in range(0, len(img), w):
ascii_str += ("".join(chars[p * scale // 255] * 2 for p in img[i : i + w])) + "\n"
logging.info(f"\n{ascii_str}")
def log_multiple() -> None:
for idx in range(10):
img, label = mnist_sample(idx)
log_ascii_mnist(img)
logging.info(f"^ {label=}")
if __name__ == "__main__":
main()