# 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
```


```python
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()
```