examples/mnist_from_scratch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST from Scratch\n",
    "\n",
    "Can we train a model to recognize handwritten digits using numpy?\n",
    "\n",
    "1. Load the MNIST dataset from the web and store as NumPy arrays\n",
    "2. Train a simple model to solve MNIST using PyTorch\n",
    "3. Do the same with NumPy by implementing various ML algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "import optax\n",
    "from flax import nnx\n",
    "from tqdm import trange\n",
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "rngs = nnx.Rngs(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load mnist dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"mnist\")\n",
    "X_train = np.array([np.array(image) for image in dataset[\"train\"][\"image\"]], dtype=np.float32).reshape(-1, 784) / 255.0\n",
    "Y_train = np.array(dataset[\"train\"][\"label\"], dtype=np.int32)\n",
    "X_test = np.array([np.array(image) for image in dataset[\"test\"][\"image\"]], dtype=np.float32).reshape(-1, 784) / 255.0\n",
    "Y_test = np.array(dataset[\"test\"][\"label\"], dtype=np.int32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve with PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nnx.Module):\n",
    "    def __init__(self, *, rngs):\n",
    "        self.l1 = nnx.Linear(784, 128, rngs=rngs)\n",
    "        self.l2 = nnx.Linear(128, 10, rngs=rngs)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        x = nnx.relu(self.l1(x))\n",
    "        x = self.l2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-0.49278238  1.2711219  -0.0858871  -0.08085708 -0.4907502  -1.1703337\n",
      "   0.41073963 -0.03593045  0.57198393 -0.7578646 ]]\n"
     ]
    }
   ],
   "source": [
    "model = Net(rngs=rngs)\n",
    "y = model(jnp.ones((1, 784)))\n",
    "nnx.display(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 0.005\n",
    "momentum = 0.9\n",
    "\n",
    "optimizer = nnx.Optimizer(model, optax.adamw(learning_rate, momentum))\n",
    "metrics = nnx.MultiMetric(accuracy=nnx.metrics.Accuracy(), loss=nnx.metrics.Average(\"loss\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_fn(model, images, labels):\n",
    "    logits = model(images)\n",
    "    loss = optax.softmax_cross_entropy_with_integer_labels(logits=logits, labels=labels).mean()\n",
    "    return loss, logits\n",
    "\n",
    "@nnx.jit\n",
    "def train_step(model, optimizer, metrics, images, labels):\n",
    "    grad_fn = nnx.value_and_grad(loss_fn, has_aux=True)\n",
    "    (loss, logits), grads = grad_fn(model, images, labels)\n",
    "    metrics.update(loss=loss, logits=logits, labels=labels)\n",
    "    optimizer.update(grads)\n",
    "\n",
    "@nnx.jit\n",
    "def eval_step(model, metrics, images, labels):\n",
    "    loss, logits = loss_fn(model, images, labels)\n",
    "    metrics.update(loss=loss, logits=logits, labels=labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[train] step: 200, loss: 0.4570, accuracy: 85.70\n",
      "[test] step: 200, loss: nan, accuracy: 91.18\n",
      "[train] step: 400, loss: 0.2540, accuracy: 92.36\n",
      "[test] step: 400, loss: nan, accuracy: 93.87\n",
      "[train] step: 600, loss: 0.1826, accuracy: 94.78\n",
      "[test] step: 600, loss: nan, accuracy: 93.11\n",
      "[train] step: 800, loss: 0.1789, accuracy: 94.75\n",
      "[test] step: 800, loss: nan, accuracy: 95.06\n",
      "[train] step: 1000, loss: 0.1614, accuracy: 95.23\n",
      "[test] step: 1000, loss: nan, accuracy: 94.14\n",
      "[train] step: 1200, loss: 0.1473, accuracy: 95.55\n",
      "[test] step: 1200, loss: nan, accuracy: 95.54\n",
      "[train] step: 1400, loss: 0.1330, accuracy: 96.05\n",
      "[test] step: 1400, loss: nan, accuracy: 95.87\n",
      "[train] step: 1600, loss: 0.1256, accuracy: 96.05\n",
      "[test] step: 1600, loss: nan, accuracy: 96.02\n",
      "[train] step: 1800, loss: 0.1266, accuracy: 95.94\n",
      "[test] step: 1800, loss: nan, accuracy: 96.11\n",
      "[train] step: 1874, loss: 0.1144, accuracy: 96.28\n",
      "[test] step: 1874, loss: nan, accuracy: 96.24\n"
     ]
    }
   ],
   "source": [
    "batch_size = 32\n",
    "eval_every = 200\n",
    "train_steps = len(X_train) // batch_size\n",
    "test_steps = len(X_test) // batch_size\n",
    "\n",
    "metrics_history = {\n",
    "    \"train_loss\": [],\n",
    "    \"train_accuracy\": [],\n",
    "    \"test_loss\": [],\n",
    "    \"test_accuracy\": [],\n",
    "}\n",
    "\n",
    "for step in range(train_steps):\n",
    "    sample = np.random.randint(0, len(X_train), size=batch_size)\n",
    "    images, labels = X_train[sample], Y_train[sample]\n",
    "    train_step(model, optimizer, metrics, images, labels)\n",
    "\n",
    "    if step > 0 and (step % eval_every == 0 or step == train_steps - 1):\n",
    "        for metric, value in metrics.compute().items():\n",
    "            metrics_history[f\"train_{metric}\"].append(value)\n",
    "        metrics.reset()\n",
    "\n",
    "        for test_step in range(test_steps):\n",
    "            images = X_test[batch_size*(test_step-1):batch_size*test_step]\n",
    "            labels = Y_test[batch_size*(test_step-1):batch_size*test_step]\n",
    "            eval_step(model, metrics, images, labels)\n",
    "\n",
    "        for metric, value in metrics.compute().items():\n",
    "            metrics_history[f\"test_{metric}\"].append(value)\n",
    "        metrics.reset()\n",
    "\n",
    "        print(\n",
    "            f\"[train] step: {step}, \"\n",
    "            f\"loss: {metrics_history['train_loss'][-1]:.4f}, \"\n",
    "            f\"accuracy: {metrics_history['train_accuracy'][-1] * 100:.2f}\"\n",
    "        )\n",
    "        print(\n",
    "            f\"[test] step: {step}, \"\n",
    "            f\"loss: {metrics_history['test_loss'][-1]:.4f}, \"\n",
    "            f\"accuracy: {metrics_history['test_accuracy'][-1] * 100:.2f}\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sanity check with NumPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "w1 = np.asarray(model.l1.kernel.value)\n",
    "w2 = np.array(model.l2.kernel.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward(x: np.ndarray) -> np.ndarray:\n",
    "    x = x.reshape(-1, 784)\n",
    "    x = x @ w1\n",
    "    x = np.maximum(x, 0)\n",
    "    x = x @ w2\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test set accuracy is 96.31\n"
     ]
    }
   ],
   "source": [
    "pred = forward(X_test).argmax(axis=1)\n",
    "accuracy = (pred == Y_test).mean() * 100\n",
    "print(f\"test set accuracy is {accuracy}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now solve with NumPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def layer_init(m: int, h: int) -> np.ndarray:\n",
    "    weights = np.random.uniform(-1., 1., size=(m, h)) / np.sqrt(m * h)\n",
    "    return weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "w1 = layer_init(784, 128)\n",
    "w2 = layer_init(128, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_forward(x0: np.ndarray) -> tuple[np.ndarray]:\n",
    "    x1 = x0 @ w1  # batch_size * 128\n",
    "    x2 = np.maximum(x1, 0)  # batch_size * 128, relu\n",
    "    x3 = x2 @ w2  # batch_size * 10\n",
    "    x3 = np.exp(x3) / np.sum(np.exp(x3), axis=1, keepdims=True)  # softmax\n",
    "    return x3, x2, x1, x0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cross_entropy_loss(pred: np.ndarray, labels: np.ndarray) -> tuple[float, np.ndarray]:\n",
    "    # NOTE: actual is the one-hot enconding of labels\n",
    "    actual = np.zeros((labels.shape[0], 10))\n",
    "    actual[np.arange(labels.shape[0]), labels] = 1\n",
    "    loss = -np.mean(np.sum(actual * np.log(pred), axis=1))\n",
    "    error = pred - actual\n",
    "    return loss, error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def backward(error: np.ndarray, xs: tuple[np.ndarray]) -> tuple[np.ndarray]:\n",
    "    _, x2, x1, x0 = xs\n",
    "    dx3 = error  # batch_size * 10, derivative of squared_error\n",
    "    dw2 = x2.T @ dx3  # 128 * 10, derivative of dot\n",
    "    dx2 = dx3 @ w2.T  # batch_size * 128, derivative of dot\n",
    "    dx1 = (x1 > 0).astype(np.float64) * dx2  # 128 * batch_size\n",
    "    dw1 = x0.T @ dx1  # 784 * 128\n",
    "    assert dw2.shape == w2.shape\n",
    "    assert dw1.shape == w1.shape\n",
    "    return dw2, dw1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_weights(dws: tuple[np.ndarray], lr: float = 1e-3) -> tuple[np.ndarray]:\n",
    "    global w1, w2\n",
    "    ws = (w2, w1)\n",
    "    for wi, dwi in zip(ws, dws):\n",
    "        wi -= lr * dwi\n",
    "    return ws"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss 0.17: 100%|██████████| 1000/1000 [00:02<00:00, 482.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test set accuracy is 94.98\n"
     ]
    }
   ],
   "source": [
    "for i in (t := trange(1000)):\n",
    "    sample = np.random.randint(0, len(X_train), size=256)\n",
    "    xs = train_forward(X_train[sample])\n",
    "    loss, error = cross_entropy_loss(xs[0], Y_train[sample].astype(int))\n",
    "    dws = backward(error, xs)\n",
    "    update_weights(dws, lr=0.001)\n",
    "    t.set_description(f\"loss {loss.item():.2f}\")\n",
    "pred = forward(X_test).argmax(axis=1)\n",
    "accuracy = (pred == Y_test).mean() * 100\n",
    "print(f\"test set accuracy is {accuracy}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}