implementation-clf-dense-mnist.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f1261642",
   "metadata": {},
   "source": [
    "## Dense classifier for classification on MNIST dataset\n",
    "\n",
    "This notebook trains a classifier for classifying the images of the classic MNIST handwritten digits dataset. This is implemented in an incremental manner, starting from a naive approach with bad performance. Every implementation addresses some issue in the previous in hope of improving the performance, and by the final version, a solid performance is reached\n",
    "<br>\n",
    "The link to the dataset used can be found in the project's README."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01089955",
   "metadata": {},
   "source": [
    "## Imports and Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "880aecf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "from ml.layers.linear import Linear\n",
    "from ml.layers.activations import Sigmoid, ReLU\n",
    "from ml.loss import BCELoss, MSELoss\n",
    "from ml.optimizers import SGD"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74c4b32b",
   "metadata": {},
   "source": [
    "### Loading data and creating data utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a78a7208",
   "metadata": {},
   "outputs": [],
   "source": [
    "# MNIST Dataset Class\n",
    "class Dataset:\n",
    "    def __init__(self, path):\n",
    "        df = pd.read_csv(path, header=None)\n",
    "        self.X = df.iloc[:, 1:]\n",
    "        self.y = df.iloc[:, 0]\n",
    "\n",
    "        self.X = self.X / 255\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.X)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return np.array(self.X.iloc[idx]), self.y[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "254af3ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(60000, 10000)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data = Dataset(\"./datasets/mnist/mnist_train.csv\")\n",
    "test_data = Dataset(\"./datasets/mnist/mnist_test.csv\")\n",
    "len(train_data), len(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f0790bf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataloader\n",
    "def get_batch_loader(dataset, batchsize):\n",
    "    n = len(dataset)\n",
    "    indices = list(range(n))\n",
    "    np.random.shuffle(indices)\n",
    "    start = 0\n",
    "    end = batchsize\n",
    "    while True:\n",
    "        if start == n:\n",
    "            start = 0\n",
    "            end = batchsize\n",
    "            np.random.shuffle(indices)\n",
    "            yield None\n",
    "            continue\n",
    "        batch_X, batch_y = list(), list()\n",
    "        for i in range(start, end):\n",
    "            X, y = train_data[i]\n",
    "            batch_X.append(X)\n",
    "            batch_y.append(y)\n",
    "        yield np.array(batch_X), np.array(batch_y)\n",
    "        start = end\n",
    "        end = min(n, end + batchsize)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed5dbb03",
   "metadata": {},
   "source": [
    "### Training and testing utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0e1cec74",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training utility\n",
    "def train_model(model, loss_fn, dataloader, optimizer, n_epochs, grad_debug=False):\n",
    "    for epoch in range(n_epochs):\n",
    "        batch_count = 0\n",
    "        grad_hist = list()\n",
    "        while True:\n",
    "            batch = next(dataloader)\n",
    "            if not batch:\n",
    "                break\n",
    "            batch_count += 1\n",
    "            x, y = batch\n",
    "            y = np.eye(10)[y]\n",
    "            y_pred = model.forward(x)\n",
    "            loss = loss_fn.forward(y, y_pred)\n",
    "            din_loss = loss_fn.backward()\n",
    "            model.backward(din_loss)\n",
    "            model.update(optimizer)\n",
    "            if grad_debug:\n",
    "                step_gradient = list()\n",
    "                for l in model.layers:\n",
    "                    if l.type == \"learnable\":\n",
    "                        step_gradient.append(l.grad)\n",
    "                grad_hist.append(step_gradient)\n",
    "            print(\n",
    "                \"\\rEpoch {}/{} Batch {} Loss {}\".format(\n",
    "                    epoch + 1, n_epochs, batch_count, loss\n",
    "                ),\n",
    "                end=\"\",\n",
    "            )\n",
    "            if np.isnan(loss):\n",
    "                raise ValueError(\"Loss value has reached NaN, something is wrong\")\n",
    "            if np.isinf(loss):\n",
    "                raise ValueError(\"Loss value has reached inf, something is wrong\")\n",
    "        print()\n",
    "        if grad_debug:\n",
    "            return grad_hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ad8377ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Testing utility\n",
    "def test_model(model, dataloader):\n",
    "    true = list()\n",
    "    pred = list()\n",
    "    while True:\n",
    "        batch = next(dataloader)\n",
    "        if not batch:\n",
    "            break\n",
    "        x, y = batch\n",
    "        y_pred = np.argmax(clf.forward(x), axis=-1)\n",
    "        true.extend(list(np.squeeze(y)))\n",
    "        pred.extend(list(np.squeeze(y_pred)))\n",
    "    print(\"Accuracy: {:.03f}%\".format(100 * accuracy_score(true, pred)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f5e4368",
   "metadata": {},
   "source": [
    "### Model Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "67028cbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Base model class\n",
    "class Model:\n",
    "    def __init__(self, layers=[]):\n",
    "        self.layers = layers\n",
    "\n",
    "    def add_layer(self, layer):\n",
    "        self.layers.append(layer)\n",
    "\n",
    "    def forward(self, x):\n",
    "        for l in self.layers:\n",
    "            x = l.forward(x)\n",
    "        return x\n",
    "\n",
    "    def backward(self, dx):\n",
    "        for l in self.layers[::-1]:\n",
    "            dx = l.backward(dx)\n",
    "\n",
    "    def update(self, optim):\n",
    "        for l in self.layers:\n",
    "            if l.type == \"learnable\":\n",
    "                for p in l.params:\n",
    "                    l.params[p] = optim.update(l.params[p], l.grad[p])\n",
    "\n",
    "    def num_params(self):\n",
    "        count = 0\n",
    "        for l in self.layers:\n",
    "            if l.type == \"learnable\":\n",
    "                for name, params in l.params.items():\n",
    "                    count += np.prod(params.shape)\n",
    "        return count"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f3d46f6",
   "metadata": {},
   "source": [
    "## Implementation 1\n",
    "- 3 dense layers followed by sigmoid\n",
    "- MSE Loss\n",
    "- SGD Optimizer with LR 0.1\n",
    "- 5 epochs\n",
    "\n",
    "A small network trained for a short duration to see how the learning goes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3619e045",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of params: 80730\n"
     ]
    }
   ],
   "source": [
    "layers = [\n",
    "    Linear(784, 100),\n",
    "    Sigmoid(),\n",
    "    Linear(100, 20),\n",
    "    Sigmoid(),\n",
    "    Linear(20, 10),\n",
    "    Sigmoid(),\n",
    "]\n",
    "clf = Model(layers)\n",
    "print(\"Number of params:\", clf.num_params())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0cbe3f4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCHSIZE = 32\n",
    "train_loader = get_batch_loader(train_data, BATCHSIZE)\n",
    "test_loader = get_batch_loader(test_data, BATCHSIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ba59d615",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "016cf13e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sgd_optim = SGD(lr=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "08b24949",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5 Batch 1875 Loss 0.047137227423678255\n",
      "Epoch 2/5 Batch 1875 Loss 0.046330285603477725\n",
      "Epoch 3/5 Batch 1875 Loss 0.046136397428807714\n",
      "Epoch 4/5 Batch 1875 Loss 0.046091304782153106\n",
      "Epoch 5/5 Batch 1875 Loss 0.046030348715061045\n"
     ]
    }
   ],
   "source": [
    "train_model(clf, loss_fn, train_loader, sgd_optim, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "281cfda2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 21.343%\n"
     ]
    }
   ],
   "source": [
    "test_model(clf, train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "babeb287",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 21.840%\n"
     ]
    }
   ],
   "source": [
    "test_model(clf, test_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac6b555e",
   "metadata": {},
   "source": [
    "This small model performs very poorly, but then it was trained for a very short duration. The next implementation trains the same network for longer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59ccc800",
   "metadata": {},
   "source": [
    "## Implementation 2\n",
    "\n",
    "- 3 dense layers followed by sigmoid\n",
    "- MSE Loss\n",
    "- SGD Optimizer with LR 0.1\n",
    "- 10 epochs\n",
    "\n",
    "This version trains the model for 10 epochs vs 5 epochs in previous version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "515628e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of params: 80730\n"
     ]
    }
   ],
   "source": [
    "layers = [\n",
    "    Linear(784, 100),\n",
    "    Sigmoid(),\n",
    "    Linear(100, 20),\n",
    "    Sigmoid(),\n",
    "    Linear(20, 10),\n",
    "    Sigmoid(),\n",
    "]\n",
    "clf = Model(layers)\n",
    "print(\"Number of params:\", clf.num_params())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "63bc46be",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCHSIZE = 32\n",
    "train_loader = get_batch_loader(train_data, BATCHSIZE)\n",
    "test_loader = get_batch_loader(test_data, BATCHSIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "60827004",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3ec5b83b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sgd_optim = SGD(lr=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "88b80435",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10 Batch 1875 Loss 0.07366473338410287\n",
      "Epoch 2/10 Batch 1875 Loss 0.045338576624558556\n",
      "Epoch 3/10 Batch 1875 Loss 0.044649676585136894\n",
      "Epoch 4/10 Batch 1875 Loss 0.043895667895039875\n",
      "Epoch 5/10 Batch 1875 Loss 0.043126566307594885\n",
      "Epoch 6/10 Batch 1875 Loss 0.042386199181076694\n",
      "Epoch 7/10 Batch 1875 Loss 0.041687917798274035\n",
      "Epoch 8/10 Batch 1875 Loss 0.041029392827415854\n",
      "Epoch 9/10 Batch 1875 Loss 0.040411216528821614\n",
      "Epoch 10/10 Batch 1875 Loss 0.039836226156774685\n"
     ]
    }
   ],
   "source": [
    "train_model(clf, loss_fn, train_loader, sgd_optim, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3c41b27d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 31.860%\n"
     ]
    }
   ],
   "source": [
    "test_model(clf, train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "472110f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 32.650%\n"
     ]
    }
   ],
   "source": [
    "test_model(clf, test_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c474b238",
   "metadata": {},
   "source": [
    "This shows minor improvement but the performance is still poor. Something needs to change.\n",
    "<br>\n",
    "The first thing to change here is the loss function. MSE loss is not a suitable loss function for the task of classification. When paired with sigmoid activation, MSE loss can cause a slowdown in learning. In general, a cross-entropy loss is used. The next implementation replaces MSE with a BCE loss (binary cross-entropy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41bd8c64",
   "metadata": {},
   "source": [
    "## Implementation 3\n",
    "- 3 dense layers followed by sigmoid\n",
    "- BCE Loss\n",
    "- SGD Optimizer with LR 0.1\n",
    "- 5 epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "54438530",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of params: 80730\n"
     ]
    }
   ],
   "source": [
    "layers = [\n",
    "    Linear(784, 100),\n",
    "    Sigmoid(),\n",
    "    Linear(100, 20),\n",
    "    Sigmoid(),\n",
    "    Linear(20, 10),\n",
    "    Sigmoid(),\n",
    "]\n",
    "clf = Model(layers)\n",
    "print(\"Number of params:\", clf.num_params())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "9657b870",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCHSIZE = 32\n",
    "train_loader = get_batch_loader(train_data, BATCHSIZE)\n",
    "test_loader = get_batch_loader(test_data, BATCHSIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "f05252e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = BCELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "bff73407",
   "metadata": {},
   "outputs": [],
   "source": [
    "sgd_optim = SGD(lr=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "074ce73d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10 Batch 1875 Loss 0.31522753038840196\n",
      "Epoch 2/10 Batch 1875 Loss 0.29131345847025495\n",
      "Epoch 3/10 Batch 1875 Loss 0.27427505168658585\n",
      "Epoch 4/10 Batch 1875 Loss 0.25775393224504945\n",
      "Epoch 5/10 Batch 1875 Loss 0.24097654013478956\n",
      "Epoch 6/10 Batch 1875 Loss 0.22670178875852232\n",
      "Epoch 7/10 Batch 1875 Loss 0.21343617909173368\n",
      "Epoch 8/10 Batch 1875 Loss 0.20111324669733924\n",
      "Epoch 9/10 Batch 1875 Loss 0.19074996649335482\n",
      "Epoch 10/10 Batch 1875 Loss 0.18200842995863412\n"
     ]
    }
   ],
   "source": [
    "train_model(clf, loss_fn, train_loader, sgd_optim, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "26a4cd3a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 68.822%\n"
     ]
    }
   ],
   "source": [
    "test_model(clf, train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "30e144ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 69.260%\n"
     ]
    }
   ],
   "source": [
    "test_model(clf, test_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a409d53",
   "metadata": {},
   "source": [
    "This is a very solid improvement. But still, there's some gap to cover. As the next step, the model is made larger and more powerful. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73a6a36d",
   "metadata": {},
   "source": [
    "## Implementation 4\n",
    "- 10 dense layers followed by sigmoid\n",
    "- BCE Loss\n",
    "- SGD Optimizer with LR 0.1\n",
    "- 10 epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "b033d77b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of params: 640740\n"
     ]
    }
   ],
   "source": [
    "layers = [\n",
    "    Linear(784, 500),\n",
    "    Sigmoid(),\n",
    "    Linear(500, 300),\n",
    "    Sigmoid(),\n",
    "    Linear(300, 200),\n",
    "    Sigmoid(),\n",
    "    Linear(200, 100),\n",
    "    Sigmoid(),\n",
    "    Linear(100, 80),\n",
    "    Sigmoid(),\n",
    "    Linear(80, 60),\n",
    "    Sigmoid(),\n",
    "    Linear(60, 50),\n",
    "    Sigmoid(),\n",
    "    Linear(50, 20),\n",
    "    Sigmoid(),\n",
    "    Linear(20, 20),\n",
    "    Sigmoid(),\n",
    "    Linear(20, 10),\n",
    "    Sigmoid(),\n",
    "]\n",
    "clf = Model(layers)\n",
    "print(\"Number of params:\", clf.num_params())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dc75bb8",
   "metadata": {},
   "source": [
    "This number of parameters increases 8x from 80k to 640k."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "279321f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCHSIZE = 32\n",
    "train_loader = get_batch_loader(train_data, BATCHSIZE)\n",
    "test_loader = get_batch_loader(test_data, BATCHSIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "4963510a",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = BCELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "8e8f487e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sgd_optim = SGD(lr=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "6eb03fe1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10 Batch 1875 Loss 0.32065267100320665\n",
      "Epoch 2/10 Batch 1875 Loss 0.30359540641930194\n",
      "Epoch 3/10 Batch 1875 Loss 0.27758699601144066\n",
      "Epoch 4/10 Batch 1875 Loss 0.24886533162364738\n",
      "Epoch 5/10 Batch 1875 Loss 0.22131867472329345\n",
      "Epoch 6/10 Batch 1875 Loss 0.19981650587817744\n",
      "Epoch 7/10 Batch 1875 Loss 0.18165645233607336\n",
      "Epoch 8/10 Batch 1875 Loss 0.16385136628219285\n",
      "Epoch 9/10 Batch 1875 Loss 0.14793189293076198\n",
      "Epoch 10/10 Batch 1875 Loss 0.13647412177745427\n"
     ]
    }
   ],
   "source": [
    "train_model(clf, loss_fn, train_loader, sgd_optim, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "f331f98e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 68.782%\n"
     ]
    }
   ],
   "source": [
    "test_model(clf, train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "bb89decd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 68.870%\n"
     ]
    }
   ],
   "source": [
    "test_model(clf, test_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5011950",
   "metadata": {},
   "source": [
    "The scores did not improve. This is counterintuitive since larger networks should learn better. To understand the issue, the behavior of the gradients can be checked."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "1c6dac5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "clf = Model(layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "98a335ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/1 Batch 1875 Loss 0.12595507255819066\n"
     ]
    }
   ],
   "source": [
    "grad_history = train_model(clf, loss_fn, train_loader, sgd_optim, 1, grad_debug=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "f468b816",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_norms = list()\n",
    "b_norms = list()\n",
    "for step_grad in grad_history:\n",
    "    for layer_grad in step_grad:\n",
    "        w_norms.append(np.linalg.norm(layer_grad[\"w\"]))\n",
    "        b_norms.append(np.linalg.norm(layer_grad[\"b\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "63219867",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean of grad-norm: 0.061534779654178674 0.008031953525073295\n",
      "Variance of grad-norm: 0.0014953861889968982 1.254105387885096e-05\n"
     ]
    }
   ],
   "source": [
    "print(\"Mean of grad-norm:\", np.mean(w_norms), np.mean(b_norms))\n",
    "print(\"Variance of grad-norm:\", np.var(w_norms), np.var(b_norms))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc27f944",
   "metadata": {},
   "source": [
    "For the analysis. the gradients for all the parameters are stored for all the steps in the first epoch. The norm is taken for the gradients which reflects some measure of size. The mean and variance of the norms for w and b parameters are shown.\n",
    "<br>\n",
    "The numbers indicate that the gradients are small hinting at the problem of vanishing gradients. This arises in deeper networks with sigmoid activations, which tend to saturate making it slow and difficult to learn. The suggested activation function to tackle vanishing gradients is ReLU. The next implementation replaces all the sigmoid activations (except final) with ReLU."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d5824b8",
   "metadata": {},
   "source": [
    "## Implementation 5\n",
    "- 10 dense layers, the first 9 followed by ReLU, last by Sigmoid\n",
    "- BCE Loss\n",
    "- SGD Optimizer with LR 0.1\n",
    "- 10 epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "425c14a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of params: 640740\n"
     ]
    }
   ],
   "source": [
    "layers = [\n",
    "    Linear(784, 500),\n",
    "    ReLU(),\n",
    "    Linear(500, 300),\n",
    "    ReLU(),\n",
    "    Linear(300, 200),\n",
    "    ReLU(),\n",
    "    Linear(200, 100),\n",
    "    ReLU(),\n",
    "    Linear(100, 80),\n",
    "    ReLU(),\n",
    "    Linear(80, 60),\n",
    "    ReLU(),\n",
    "    Linear(60, 50),\n",
    "    ReLU(),\n",
    "    Linear(50, 20),\n",
    "    ReLU(),\n",
    "    Linear(20, 20),\n",
    "    ReLU(),\n",
    "    Linear(20, 10),\n",
    "    Sigmoid(),\n",
    "]\n",
    "clf = Model(layers)\n",
    "print(\"Number of params:\", clf.num_params())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "8772128d",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCHSIZE = 32\n",
    "train_loader = get_batch_loader(train_data, BATCHSIZE)\n",
    "test_loader = get_batch_loader(test_data, BATCHSIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "c8398321",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = BCELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "554a5bf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sgd_optim = SGD(lr=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "96c0d786",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r",
      "Epoch 1/1 Batch 1 Loss inf"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/himansu/Programming/general_ml/ml-from-scratch/ml/layers/activations.py:24: RuntimeWarning: overflow encountered in exp\n",
      "  return np.where(x >= 0, 1 / (1 + np.nan_to_num(np.exp(-x))), np.nan_to_num(np.exp(x)) / (1 + np.nan_to_num(np.exp(x))))\n",
      "/home/himansu/Programming/general_ml/ml-from-scratch/ml/loss.py:40: RuntimeWarning: divide by zero encountered in log\n",
      "  sample_loss = -(true * np.nan_to_num(np.log(pred)) + (1 - true) * np.nan_to_num(np.log(1 - pred)))\n",
      "/home/himansu/.virtualenvs/general-ml-env/lib/python3.10/site-packages/numpy/core/_methods.py:179: RuntimeWarning: overflow encountered in reduce\n",
      "  ret = umr_sum(arr, axis, dtype, out, keepdims, where=where)\n",
      "/home/himansu/Programming/general_ml/ml-from-scratch/ml/loss.py:46: RuntimeWarning: divide by zero encountered in true_divide\n",
      "  return -(true / pred - (1 - true) / (1 - pred)) / (b * n)\n",
      "/home/himansu/Programming/general_ml/ml-from-scratch/ml/loss.py:46: RuntimeWarning: invalid value encountered in true_divide\n",
      "  return -(true / pred - (1 - true) / (1 - pred)) / (b * n)\n",
      "/home/himansu/Programming/general_ml/ml-from-scratch/ml/layers/activations.py:21: RuntimeWarning: invalid value encountered in multiply\n",
      "  return sigmoid_x * (1 - sigmoid_x) * dy\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Loss value has reached inf, something is wrong",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Input \u001b[0;32mIn [44]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m grad_history \u001b[38;5;241m=\u001b[39m \u001b[43mtrain_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43mclf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mloss_fn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msgd_optim\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgrad_debug\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n",
      "Input \u001b[0;32mIn [5]\u001b[0m, in \u001b[0;36mtrain_model\u001b[0;34m(model, loss_fn, dataloader, optimizer, n_epochs, grad_debug)\u001b[0m\n\u001b[1;32m     28\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mLoss value has reached NaN, something is wrong\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     29\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m np\u001b[38;5;241m.\u001b[39misinf(loss):\n\u001b[0;32m---> 30\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mLoss value has reached inf, something is wrong\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     31\u001b[0m \u001b[38;5;28mprint\u001b[39m()\n\u001b[1;32m     32\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m grad_debug:\n",
      "\u001b[0;31mValueError\u001b[0m: Loss value has reached inf, something is wrong"
     ]
    }
   ],
   "source": [
    "grad_history = train_model(clf, loss_fn, train_loader, sgd_optim, 1, grad_debug=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55417d9f",
   "metadata": {},
   "source": [
    "The loss value reaches infinity in the first iteration itself. To check what’s wrong, the model is initialized and a single step of forward-pass is taken."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "e3d25a17",
   "metadata": {},
   "outputs": [],
   "source": [
    "clf_test = Model(layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "84c0558d",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = next(train_loader)\n",
    "y = np.eye(10)[y]\n",
    "y_pred = clf.forward(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "4b2765f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "inf"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss_fn.forward(y, y_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "69d836f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cb8a2e1",
   "metadata": {},
   "source": [
    "All the predictions are zero. This happens with sigmoid activations when the input is very negative. This behavior can be caused by improper weight initialization. By default, the weights in the Linear layer are initialized from a unit normal distribution. The below implementation decreases the variance of the weights to 0.1.\n",
    "<br>\n",
    "But before that, the question is why did this happen with ReLU and didn’t happen with sigmoid? The outputs from the sigmoid are bounded to a range [0, 1]. So even if weights are large, the overall weighted sum is still small. When switching to ReLU, the unbounded nature of input changes to unbounded, which can result in very large activations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "1475d834",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of params: 640740\n"
     ]
    }
   ],
   "source": [
    "def small_weight_init(val_range):\n",
    "    def init(dims):\n",
    "        in_d, out_d = dims\n",
    "        return np.random.randn(in_d, out_d) * val_range\n",
    "\n",
    "    return init\n",
    "\n",
    "\n",
    "layers = [\n",
    "    Linear(784, 500, small_weight_init(0.1)),\n",
    "    ReLU(),\n",
    "    Linear(500, 300, small_weight_init(0.1)),\n",
    "    ReLU(),\n",
    "    Linear(300, 200, small_weight_init(0.1)),\n",
    "    ReLU(),\n",
    "    Linear(200, 100, small_weight_init(0.1)),\n",
    "    ReLU(),\n",
    "    Linear(100, 80, small_weight_init(0.1)),\n",
    "    ReLU(),\n",
    "    Linear(80, 60, small_weight_init(0.1)),\n",
    "    ReLU(),\n",
    "    Linear(60, 50, small_weight_init(0.1)),\n",
    "    ReLU(),\n",
    "    Linear(50, 20, small_weight_init(0.1)),\n",
    "    ReLU(),\n",
    "    Linear(20, 20, small_weight_init(0.1)),\n",
    "    ReLU(),\n",
    "    Linear(20, 10, small_weight_init(0.1)),\n",
    "    Sigmoid(),\n",
    "]\n",
    "clf = Model(layers)\n",
    "print(\"Number of params:\", clf.num_params())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "f5a9ab1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCHSIZE = 32\n",
    "train_loader = get_batch_loader(train_data, BATCHSIZE)\n",
    "test_loader = get_batch_loader(test_data, BATCHSIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "9032424c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/1 Batch 1875 Loss 0.14808825166565814\n"
     ]
    }
   ],
   "source": [
    "grad_history = train_model(clf, loss_fn, train_loader, sgd_optim, 1, grad_debug=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dad1938",
   "metadata": {},
   "source": [
    "With smaller weights, no instability was faced and the model trained without any issues. Now, we can check the behaviour of the gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "076534df",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_norms = list()\n",
    "b_norms = list()\n",
    "for step_grad in grad_history:\n",
    "    for layer_grad in step_grad:\n",
    "        w_norms.append(np.linalg.norm(layer_grad[\"w\"]))\n",
    "        b_norms.append(np.linalg.norm(layer_grad[\"b\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "289e48a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean of grad-norm: 0.15127703638150822 0.008138200530934742\n",
      "Variance of grad-norm: 0.002460627057115184 3.668307440557543e-05\n"
     ]
    }
   ],
   "source": [
    "print(\"Mean of grad-norm:\", np.mean(w_norms), np.mean(b_norms))\n",
    "print(\"Variance of grad-norm:\", np.var(w_norms), np.var(b_norms))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c7d2ccc",
   "metadata": {},
   "source": [
    "The means of gradient norms are now much larger. This suggests that the learning slowdown shouldn't be happening here. This version is trained for 9 further epochs (since 1 epoch is already trained)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "7b6d772e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/9 Batch 1875 Loss 0.0164925605651016426\n",
      "Epoch 2/9 Batch 1875 Loss 0.0069168688169994466\n",
      "Epoch 3/9 Batch 1875 Loss 0.0020783844567613804\n",
      "Epoch 4/9 Batch 1875 Loss 0.00151930843760857617\n",
      "Epoch 5/9 Batch 1875 Loss 0.00108543128782836725\n",
      "Epoch 6/9 Batch 1875 Loss 0.00129777724346944255\n",
      "Epoch 7/9 Batch 1875 Loss 0.00334526778004999451\n",
      "Epoch 8/9 Batch 1875 Loss 0.00120160082590837105\n",
      "Epoch 9/9 Batch 1875 Loss 0.00076900947263972114\n"
     ]
    }
   ],
   "source": [
    "train_model(clf, loss_fn, train_loader, sgd_optim, 9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "9d9bfdc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 99.002%\n"
     ]
    }
   ],
   "source": [
    "test_model(clf, train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "1ff8e181",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 99.000%\n"
     ]
    }
   ],
   "source": [
    "test_model(clf, test_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d66c1935",
   "metadata": {},
   "source": [
    "This version crosses 99% accuracy- a very huge improvement. The first implementation started from accuracy in the 20s, and after incremental improvements, has reached almost perfect accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5be94ac1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "general-ml-env",
   "language": "python",
   "name": "general-ml-env"
  },
  "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}