pytorch-vision/non-linear-mnist.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-layer non-linear network for classifying handwritten digits\n",
    "\n",
    "In this notebook I will be using one of MNIST datasets to train a multi-layer non-linear network to classify handwritten digits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, torchvision\n",
    "from torch import nn\n",
    "from torchvision import datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fetch the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = datasets.MNIST(\n",
    "    root=\".datasets\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=torchvision.transforms.ToTensor(),\n",
    "    target_transform=None\n",
    ")\n",
    "\n",
    "test_data = datasets.MNIST(\n",
    "    root=\".datasets\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=torchvision.transforms.ToTensor(),\n",
    "    target_transform=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<torch.utils.data.dataloader.DataLoader at 0x7ecc64a7a8a0>,\n",
       " <torch.utils.data.dataloader.DataLoader at 0x7ecc64a2a270>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    dataset=train_data,\n",
    "    batch_size=32,\n",
    "    shuffle=True\n",
    ")\n",
    "\n",
    "test_dataloader = torch.utils.data.DataLoader(\n",
    "    dataset=test_data,\n",
    "    batch_size=32,\n",
    "    shuffle=False\n",
    ")\n",
    "\n",
    "train_dataloader, test_dataloader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MNISTModelV0(nn.Module):\n",
    "    def __init__(self, in_shape, hidden_units, out_shape):\n",
    "        super().__init__()\n",
    "        self.layer_stack = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(in_features=in_shape, out_features=hidden_units),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(in_features=hidden_units, out_features=hidden_units),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(in_features=hidden_units, out_features=hidden_units),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(in_features=hidden_units, out_features=out_shape)\n",
    "        )\n",
    "    \n",
    "    def forward(self, X):\n",
    "        return self.layer_stack.forward(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(64)\n",
    "\n",
    "model = MNISTModelV0(\n",
    "    in_shape=28*28,\n",
    "    hidden_units=10,\n",
    "    out_shape=len(train_data.classes)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training and testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 56.25% | Loss: 0.7739577293395996\n",
      "Accuracy: 81.25% | Loss: 0.4482867121696472\n",
      "Accuracy: 81.25% | Loss: 0.3499242663383484\n",
      "Accuracy: 87.50% | Loss: 0.20318855345249176\n",
      "Accuracy: 81.25% | Loss: 0.26251253485679626\n",
      "Accuracy: 87.50% | Loss: 0.23953774571418762\n",
      "Accuracy: 87.50% | Loss: 0.2038857638835907\n",
      "Accuracy: 93.75% | Loss: 0.24988947808742523\n",
      "Accuracy: 100.00% | Loss: 0.09009826183319092\n",
      "Accuracy: 93.75% | Loss: 0.1330987960100174\n",
      "Accuracy: 93.75% | Loss: 0.15335838496685028\n",
      "Accuracy: 93.75% | Loss: 0.21128623187541962\n",
      "Accuracy: 93.75% | Loss: 0.11609114706516266\n",
      "Accuracy: 100.00% | Loss: 0.04910185560584068\n",
      "Accuracy: 93.75% | Loss: 0.06536896526813507\n",
      "Accuracy: 93.75% | Loss: 0.08226760476827621\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(params=model.parameters(), lr=0.01)\n",
    "epochs = 16\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    for batch, (X, y) in enumerate(train_dataloader):\n",
    "        model.train()\n",
    "        y_logits = model.forward(X)\n",
    "        loss = loss_fn(y_logits, y)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "    for batch, (X, y) in enumerate(test_dataloader):\n",
    "        model.eval()\n",
    "        with torch.inference_mode():\n",
    "            y_logits = model.forward(X)\n",
    "            y_pred = y_logits.argmax(dim=1)\n",
    "            loss = loss_fn(y_logits, y)\n",
    "        accuracy = accuracy_score(y, y_pred)\n",
    "    \n",
    "    print(f\"Accuracy: {accuracy*100:.2f}% | Loss: {loss}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even though it managed to hit `100%` on some iterations, all by chance (I think), it is still not any better than the linear model. Both got `93.75%` accuracy. Oh well..."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}