MNIST_reference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -q  torch  torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleMLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleMLP, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.layer_stack = nn.Sequential(\n",
    "            nn.Linear(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 10),\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.layer_stack(x)\n",
    "        return logits\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transformations applied on each image\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,))\n",
    "])\n",
    "\n",
    "# Loading MNIST dataset from torchvision\n",
    "train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)\n",
    "\n",
    "# # Data Loaders\n",
    "# train_loader = DataLoader(train_dataset, batch_size=64, shuffle=False)\n",
    "# test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)\n",
    "\n",
    "from torch.utils.data import Subset\n",
    "\n",
    "# Sort the dataset by labels\n",
    "def sort_dataset_by_labels(dataset):\n",
    "    indices = list(range(len(dataset)))\n",
    "    indices.sort(key=lambda x: dataset.targets[x])\n",
    "    return Subset(dataset, indices)\n",
    "\n",
    "sorted_train_dataset = sort_dataset_by_labels(train_dataset)\n",
    "sorted_test_dataset = sort_dataset_by_labels(test_dataset)\n",
    "\n",
    "train_loader = DataLoader(sorted_train_dataset, batch_size=64, shuffle=False)\n",
    "test_loader = DataLoader(sorted_test_dataset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Labels for the first 100 items: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n"
     ]
    }
   ],
   "source": [
    "labels_list = []\n",
    "\n",
    "# Iterate through the DataLoader\n",
    "for images, labels in train_loader:\n",
    "    labels_list.extend(labels.tolist())  # Convert labels to a Python list and add to the labels_list\n",
    "    if len(labels_list) >= 100:  # Check if we have collected 100 labels\n",
    "        break\n",
    "\n",
    "# Print the first 100 labels\n",
    "print(\"Labels for the first 100 items:\", labels_list[:100])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'classic': 668672, 'bio': 553600.0}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# number of parameters in the classic vs bio model\n",
    "{'classic': 784*512 + 512*512 + 512*10, 'bio': (768*2000 + 2000*2000) / 10}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'classic': 668672, 'bio': 276800.0}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "{'classic': 784*512 + 512*512 + 512*10, 'bio': 0.5 * (768*2000 + 2000*2000) / 10}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SimpleMLP(\n",
       "  (flatten): Flatten(start_dim=1, end_dim=-1)\n",
       "  (layer_stack): Sequential(\n",
       "    (0): Linear(in_features=784, out_features=512, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): Linear(in_features=512, out_features=512, bias=True)\n",
       "    (3): ReLU()\n",
       "    (4): Linear(in_features=512, out_features=10, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Model, Loss, and Optimizer\n",
    "model = SimpleMLP()\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# Device configuration\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1\n",
      "-------------------------------\n",
      "loss: 5.933360  [    0/60000]\n",
      "loss: 5.528784  [  640/60000]\n",
      "loss: 4.458291  [ 1280/60000]\n",
      "loss: 2.067710  [ 1920/60000]\n",
      "loss: 0.340005  [ 2560/60000]\n",
      "loss: 0.070007  [ 3200/60000]\n",
      "loss: 0.192108  [ 3840/60000]\n",
      "loss: 0.000535  [ 4480/60000]\n",
      "loss: 0.003097  [ 5120/60000]\n",
      "loss: 0.000879  [ 5760/60000]\n",
      "loss: 2.957772  [ 6400/60000]\n",
      "loss: 0.981879  [ 7040/60000]\n",
      "loss: 0.066049  [ 7680/60000]\n",
      "loss: 0.002220  [ 8320/60000]\n",
      "loss: 0.000964  [ 8960/60000]\n",
      "loss: 0.001016  [ 9600/60000]\n",
      "loss: 0.000382  [10240/60000]\n",
      "loss: 0.000440  [10880/60000]\n",
      "loss: 0.000999  [11520/60000]\n",
      "loss: 0.000275  [12160/60000]\n",
      "loss: 7.257589  [12800/60000]\n",
      "loss: 1.045243  [13440/60000]\n",
      "loss: 0.005094  [14080/60000]\n",
      "loss: 0.000268  [14720/60000]\n",
      "loss: 0.000016  [15360/60000]\n",
      "loss: 0.000016  [16000/60000]\n",
      "loss: 0.000030  [16640/60000]\n",
      "loss: 0.000024  [17280/60000]\n",
      "loss: 0.000003  [17920/60000]\n",
      "loss: 0.074737  [18560/60000]\n",
      "loss: 3.331630  [19200/60000]\n",
      "loss: 2.298141  [19840/60000]\n",
      "loss: 0.398009  [20480/60000]\n",
      "loss: 0.018798  [21120/60000]\n",
      "loss: 0.003219  [21760/60000]\n",
      "loss: 0.001698  [22400/60000]\n",
      "loss: 0.001369  [23040/60000]\n",
      "loss: 0.000844  [23680/60000]\n",
      "loss: 0.000499  [24320/60000]\n",
      "loss: 6.587247  [24960/60000]\n",
      "loss: 2.484885  [25600/60000]\n",
      "loss: 1.372366  [26240/60000]\n",
      "loss: 0.310959  [26880/60000]\n",
      "loss: 0.009869  [27520/60000]\n",
      "loss: 0.004086  [28160/60000]\n",
      "loss: 0.016054  [28800/60000]\n",
      "loss: 0.000479  [29440/60000]\n",
      "loss: 0.014501  [30080/60000]\n",
      "loss: 10.380549  [30720/60000]\n",
      "loss: 1.769985  [31360/60000]\n",
      "loss: 1.313237  [32000/60000]\n",
      "loss: 0.876812  [32640/60000]\n",
      "loss: 0.621755  [33280/60000]\n",
      "loss: 0.463711  [33920/60000]\n",
      "loss: 0.357888  [34560/60000]\n",
      "loss: 0.285171  [35200/60000]\n",
      "loss: 0.233775  [35840/60000]\n",
      "loss: 4.533941  [36480/60000]\n",
      "loss: 0.000356  [37120/60000]\n",
      "loss: 0.000000  [37760/60000]\n",
      "loss: 0.000000  [38400/60000]\n",
      "loss: 0.000000  [39040/60000]\n",
      "loss: 0.000000  [39680/60000]\n",
      "loss: 0.000000  [40320/60000]\n",
      "loss: 0.000000  [40960/60000]\n",
      "loss: 0.000000  [41600/60000]\n",
      "loss: 4.783977  [42240/60000]\n",
      "loss: 3.865182  [42880/60000]\n",
      "loss: 2.841220  [43520/60000]\n",
      "loss: 2.070206  [44160/60000]\n",
      "loss: 1.530899  [44800/60000]\n",
      "loss: 1.104303  [45440/60000]\n",
      "loss: 0.803271  [46080/60000]\n",
      "loss: 0.527101  [46720/60000]\n",
      "loss: 0.366340  [47360/60000]\n",
      "loss: 0.263955  [48000/60000]\n",
      "loss: 3.151761  [48640/60000]\n",
      "loss: 2.389131  [49280/60000]\n",
      "loss: 1.729616  [49920/60000]\n",
      "loss: 1.276150  [50560/60000]\n",
      "loss: 1.019827  [51200/60000]\n",
      "loss: 0.806333  [51840/60000]\n",
      "loss: 0.620184  [52480/60000]\n",
      "loss: 0.508122  [53120/60000]\n",
      "loss: 0.422892  [53760/60000]\n",
      "loss: 2.499337  [54400/60000]\n",
      "loss: 1.999729  [55040/60000]\n",
      "loss: 1.503271  [55680/60000]\n",
      "loss: 1.120197  [56320/60000]\n",
      "loss: 0.859085  [56960/60000]\n",
      "loss: 0.667894  [57600/60000]\n",
      "loss: 0.537568  [58240/60000]\n",
      "loss: 0.445670  [58880/60000]\n",
      "loss: 0.377361  [59520/60000]\n",
      "Epoch 2\n",
      "-------------------------------\n",
      "loss: 4.241849  [    0/60000]\n",
      "loss: 3.766757  [  640/60000]\n",
      "loss: 1.850675  [ 1280/60000]\n",
      "loss: 0.058538  [ 1920/60000]\n",
      "loss: 0.014453  [ 2560/60000]\n",
      "loss: 0.070649  [ 3200/60000]\n",
      "loss: 0.106292  [ 3840/60000]\n",
      "loss: 0.000939  [ 4480/60000]\n",
      "loss: 0.003509  [ 5120/60000]\n",
      "loss: 0.000821  [ 5760/60000]\n",
      "loss: 2.136283  [ 6400/60000]\n",
      "loss: 0.910588  [ 7040/60000]\n",
      "loss: 0.144348  [ 7680/60000]\n",
      "loss: 0.010018  [ 8320/60000]\n",
      "loss: 0.003616  [ 8960/60000]\n",
      "loss: 0.003242  [ 9600/60000]\n",
      "loss: 0.001366  [10240/60000]\n",
      "loss: 0.001362  [10880/60000]\n",
      "loss: 0.002743  [11520/60000]\n",
      "loss: 0.000963  [12160/60000]\n",
      "loss: 9.011411  [12800/60000]\n",
      "loss: 3.683111  [13440/60000]\n",
      "loss: 0.452089  [14080/60000]\n",
      "loss: 0.008544  [14720/60000]\n",
      "loss: 0.000659  [15360/60000]\n",
      "loss: 0.000575  [16000/60000]\n",
      "loss: 0.000369  [16640/60000]\n",
      "loss: 0.000108  [17280/60000]\n",
      "loss: 0.000147  [17920/60000]\n",
      "loss: 0.103622  [18560/60000]\n",
      "loss: 0.421914  [19200/60000]\n",
      "loss: 0.335280  [19840/60000]\n",
      "loss: 0.032317  [20480/60000]\n",
      "loss: 0.004793  [21120/60000]\n",
      "loss: 0.010119  [21760/60000]\n",
      "loss: 0.000692  [22400/60000]\n",
      "loss: 0.000665  [23040/60000]\n",
      "loss: 0.000326  [23680/60000]\n",
      "loss: 0.000136  [24320/60000]\n",
      "loss: 3.311721  [24960/60000]\n",
      "loss: 2.263369  [25600/60000]\n",
      "loss: 0.121835  [26240/60000]\n",
      "loss: 0.021754  [26880/60000]\n",
      "loss: 0.000163  [27520/60000]\n",
      "loss: 0.000405  [28160/60000]\n",
      "loss: 0.001088  [28800/60000]\n",
      "loss: 0.000131  [29440/60000]\n",
      "loss: 0.002544  [30080/60000]\n",
      "loss: 14.283473  [30720/60000]\n",
      "loss: 2.076360  [31360/60000]\n",
      "loss: 1.398301  [32000/60000]\n",
      "loss: 0.840906  [32640/60000]\n",
      "loss: 0.547094  [33280/60000]\n",
      "loss: 0.395841  [33920/60000]\n",
      "loss: 0.307487  [34560/60000]\n",
      "loss: 0.249169  [35200/60000]\n",
      "loss: 0.208014  [35840/60000]\n",
      "loss: 2.792655  [36480/60000]\n",
      "loss: 0.277147  [37120/60000]\n",
      "loss: 0.002200  [37760/60000]\n",
      "loss: 0.000005  [38400/60000]\n",
      "loss: 0.005638  [39040/60000]\n",
      "loss: 0.000000  [39680/60000]\n",
      "loss: 0.000001  [40320/60000]\n",
      "loss: 0.000037  [40960/60000]\n",
      "loss: 0.000000  [41600/60000]\n",
      "loss: 3.772929  [42240/60000]\n",
      "loss: 2.914006  [42880/60000]\n",
      "loss: 2.082560  [43520/60000]\n",
      "loss: 1.528142  [44160/60000]\n",
      "loss: 0.907819  [44800/60000]\n",
      "loss: 0.226986  [45440/60000]\n",
      "loss: 0.010183  [46080/60000]\n",
      "loss: 0.003168  [46720/60000]\n",
      "loss: 0.000595  [47360/60000]\n",
      "loss: 0.000842  [48000/60000]\n",
      "loss: 5.247957  [48640/60000]\n",
      "loss: 1.488950  [49280/60000]\n",
      "loss: 1.151364  [49920/60000]\n",
      "loss: 0.888226  [50560/60000]\n",
      "loss: 0.705398  [51200/60000]\n",
      "loss: 0.570460  [51840/60000]\n",
      "loss: 0.476711  [52480/60000]\n",
      "loss: 0.400747  [53120/60000]\n",
      "loss: 0.341480  [53760/60000]\n",
      "loss: 2.488842  [54400/60000]\n",
      "loss: 1.862356  [55040/60000]\n",
      "loss: 1.265172  [55680/60000]\n",
      "loss: 0.897600  [56320/60000]\n",
      "loss: 0.639236  [56960/60000]\n",
      "loss: 0.462150  [57600/60000]\n",
      "loss: 0.365352  [58240/60000]\n",
      "loss: 0.298701  [58880/60000]\n",
      "loss: 0.251820  [59520/60000]\n",
      "Epoch 3\n",
      "-------------------------------\n",
      "loss: 5.076498  [    0/60000]\n",
      "loss: 4.897382  [  640/60000]\n",
      "loss: 3.461470  [ 1280/60000]\n",
      "loss: 1.345560  [ 1920/60000]\n",
      "loss: 0.418355  [ 2560/60000]\n",
      "loss: 0.100667  [ 3200/60000]\n",
      "loss: 0.135272  [ 3840/60000]\n",
      "loss: 0.007392  [ 4480/60000]\n",
      "loss: 0.012153  [ 5120/60000]\n",
      "loss: 0.003625  [ 5760/60000]\n",
      "loss: 1.987553  [ 6400/60000]\n",
      "loss: 1.227084  [ 7040/60000]\n",
      "loss: 0.670684  [ 7680/60000]\n",
      "loss: 0.179954  [ 8320/60000]\n",
      "loss: 0.040976  [ 8960/60000]\n",
      "loss: 0.014318  [ 9600/60000]\n",
      "loss: 0.005797  [10240/60000]\n",
      "loss: 0.004161  [10880/60000]\n",
      "loss: 0.004517  [11520/60000]\n",
      "loss: 0.002323  [12160/60000]\n",
      "loss: 7.318002  [12800/60000]\n",
      "loss: 2.240490  [13440/60000]\n",
      "loss: 0.389237  [14080/60000]\n",
      "loss: 0.033932  [14720/60000]\n",
      "loss: 0.003996  [15360/60000]\n",
      "loss: 0.001875  [16000/60000]\n",
      "loss: 0.000804  [16640/60000]\n",
      "loss: 0.000690  [17280/60000]\n",
      "loss: 0.000681  [17920/60000]\n",
      "loss: 0.047524  [18560/60000]\n",
      "loss: 1.611749  [19200/60000]\n",
      "loss: 0.743719  [19840/60000]\n",
      "loss: 0.071150  [20480/60000]\n",
      "loss: 0.015631  [21120/60000]\n",
      "loss: 0.010409  [21760/60000]\n",
      "loss: 0.005006  [22400/60000]\n",
      "loss: 0.003343  [23040/60000]\n",
      "loss: 0.002852  [23680/60000]\n",
      "loss: 0.001902  [24320/60000]\n",
      "loss: 3.740757  [24960/60000]\n",
      "loss: 1.868791  [25600/60000]\n",
      "loss: 0.112809  [26240/60000]\n",
      "loss: 0.025455  [26880/60000]\n",
      "loss: 0.002394  [27520/60000]\n",
      "loss: 0.004743  [28160/60000]\n",
      "loss: 0.006756  [28800/60000]\n",
      "loss: 0.001599  [29440/60000]\n",
      "loss: 0.010187  [30080/60000]\n",
      "loss: 7.403306  [30720/60000]\n",
      "loss: 1.352981  [31360/60000]\n",
      "loss: 0.810305  [32000/60000]\n",
      "loss: 0.232662  [32640/60000]\n",
      "loss: 0.052006  [33280/60000]\n",
      "loss: 0.009853  [33920/60000]\n",
      "loss: 0.007384  [34560/60000]\n",
      "loss: 0.004233  [35200/60000]\n",
      "loss: 0.005593  [35840/60000]\n",
      "loss: 0.320913  [36480/60000]\n",
      "loss: 0.000076  [37120/60000]\n",
      "loss: 0.000014  [37760/60000]\n",
      "loss: 0.000000  [38400/60000]\n",
      "loss: 0.000149  [39040/60000]\n",
      "loss: 0.000007  [39680/60000]\n",
      "loss: 0.000000  [40320/60000]\n",
      "loss: 0.000002  [40960/60000]\n",
      "loss: 0.000000  [41600/60000]\n",
      "loss: 2.262742  [42240/60000]\n",
      "loss: 1.850022  [42880/60000]\n",
      "loss: 0.282948  [43520/60000]\n",
      "loss: 0.010368  [44160/60000]\n",
      "loss: 0.000650  [44800/60000]\n",
      "loss: 0.000017  [45440/60000]\n",
      "loss: 0.000022  [46080/60000]\n",
      "loss: 0.004045  [46720/60000]\n",
      "loss: 0.000024  [47360/60000]\n",
      "loss: 0.000008  [48000/60000]\n",
      "loss: 2.085424  [48640/60000]\n",
      "loss: 1.307786  [49280/60000]\n",
      "loss: 1.023736  [49920/60000]\n",
      "loss: 0.880424  [50560/60000]\n",
      "loss: 0.574142  [51200/60000]\n",
      "loss: 0.433306  [51840/60000]\n",
      "loss: 0.346344  [52480/60000]\n",
      "loss: 0.288055  [53120/60000]\n",
      "loss: 0.234636  [53760/60000]\n",
      "loss: 3.096566  [54400/60000]\n",
      "loss: 1.683063  [55040/60000]\n",
      "loss: 1.338054  [55680/60000]\n",
      "loss: 0.857949  [56320/60000]\n",
      "loss: 0.521265  [56960/60000]\n",
      "loss: 0.376315  [57600/60000]\n",
      "loss: 0.295657  [58240/60000]\n",
      "loss: 0.255519  [58880/60000]\n",
      "loss: 0.204152  [59520/60000]\n",
      "Epoch 4\n",
      "-------------------------------\n",
      "loss: 5.769371  [    0/60000]\n",
      "loss: 4.843155  [  640/60000]\n",
      "loss: 4.043718  [ 1280/60000]\n",
      "loss: 2.744131  [ 1920/60000]\n",
      "loss: 1.437196  [ 2560/60000]\n",
      "loss: 0.000145  [ 3200/60000]\n",
      "loss: 0.000000  [ 3840/60000]\n",
      "loss: 0.000000  [ 4480/60000]\n",
      "loss: 0.000000  [ 5120/60000]\n",
      "loss: 0.000000  [ 5760/60000]\n",
      "loss: 0.062041  [ 6400/60000]\n",
      "loss: 0.000735  [ 7040/60000]\n",
      "loss: 0.000000  [ 7680/60000]\n",
      "loss: 0.000000  [ 8320/60000]\n",
      "loss: 0.000000  [ 8960/60000]\n",
      "loss: 0.000000  [ 9600/60000]\n",
      "loss: 0.000000  [10240/60000]\n",
      "loss: 0.000000  [10880/60000]\n",
      "loss: 0.000000  [11520/60000]\n",
      "loss: 0.000000  [12160/60000]\n",
      "loss: 21.211267  [12800/60000]\n",
      "loss: 1.542303  [13440/60000]\n",
      "loss: 0.049810  [14080/60000]\n",
      "loss: 0.001903  [14720/60000]\n",
      "loss: 0.000228  [15360/60000]\n",
      "loss: 0.000047  [16000/60000]\n",
      "loss: 0.000029  [16640/60000]\n",
      "loss: 0.000167  [17280/60000]\n",
      "loss: 0.000300  [17920/60000]\n",
      "loss: 0.051598  [18560/60000]\n",
      "loss: 2.990248  [19200/60000]\n",
      "loss: 0.481374  [19840/60000]\n",
      "loss: 0.005268  [20480/60000]\n",
      "loss: 0.000575  [21120/60000]\n",
      "loss: 0.000165  [21760/60000]\n",
      "loss: 0.000153  [22400/60000]\n",
      "loss: 0.000081  [23040/60000]\n",
      "loss: 0.000108  [23680/60000]\n",
      "loss: 0.000046  [24320/60000]\n",
      "loss: 7.313146  [24960/60000]\n",
      "loss: 3.362474  [25600/60000]\n",
      "loss: 2.374758  [26240/60000]\n",
      "loss: 1.388846  [26880/60000]\n",
      "loss: 0.836724  [27520/60000]\n",
      "loss: 0.527943  [28160/60000]\n",
      "loss: 0.355805  [28800/60000]\n",
      "loss: 0.253927  [29440/60000]\n",
      "loss: 0.187326  [30080/60000]\n",
      "loss: 5.564179  [30720/60000]\n",
      "loss: 3.868085  [31360/60000]\n",
      "loss: 2.521529  [32000/60000]\n",
      "loss: 1.503717  [32640/60000]\n",
      "loss: 0.947541  [33280/60000]\n",
      "loss: 0.632125  [33920/60000]\n",
      "loss: 0.430020  [34560/60000]\n",
      "loss: 0.312939  [35200/60000]\n",
      "loss: 0.262832  [35840/60000]\n",
      "loss: 0.820857  [36480/60000]\n",
      "loss: 0.000384  [37120/60000]\n",
      "loss: 0.000076  [37760/60000]\n",
      "loss: 0.000000  [38400/60000]\n",
      "loss: 0.000716  [39040/60000]\n",
      "loss: 0.000014  [39680/60000]\n",
      "loss: 0.000001  [40320/60000]\n",
      "loss: 0.000023  [40960/60000]\n",
      "loss: 0.000000  [41600/60000]\n",
      "loss: 4.225939  [42240/60000]\n",
      "loss: 2.099713  [42880/60000]\n",
      "loss: 0.021048  [43520/60000]\n",
      "loss: 0.000312  [44160/60000]\n",
      "loss: 0.000052  [44800/60000]\n",
      "loss: 0.000006  [45440/60000]\n",
      "loss: 0.000004  [46080/60000]\n",
      "loss: 0.000009  [46720/60000]\n",
      "loss: 0.000004  [47360/60000]\n",
      "loss: 0.000004  [48000/60000]\n",
      "loss: 1.183003  [48640/60000]\n",
      "loss: 0.770420  [49280/60000]\n",
      "loss: 0.053400  [49920/60000]\n",
      "loss: 0.107445  [50560/60000]\n",
      "loss: 0.009155  [51200/60000]\n",
      "loss: 0.000007  [51840/60000]\n",
      "loss: 0.000211  [52480/60000]\n",
      "loss: 0.000079  [53120/60000]\n",
      "loss: 0.000770  [53760/60000]\n",
      "loss: 14.169936  [54400/60000]\n",
      "loss: 1.422236  [55040/60000]\n",
      "loss: 0.833345  [55680/60000]\n",
      "loss: 0.541925  [56320/60000]\n",
      "loss: 0.368049  [56960/60000]\n",
      "loss: 0.246617  [57600/60000]\n",
      "loss: 0.197736  [58240/60000]\n",
      "loss: 0.165923  [58880/60000]\n",
      "loss: 0.143560  [59520/60000]\n",
      "Epoch 5\n",
      "-------------------------------\n",
      "loss: 4.663995  [    0/60000]\n",
      "loss: 1.638751  [  640/60000]\n",
      "loss: 0.000132  [ 1280/60000]\n",
      "loss: 0.000000  [ 1920/60000]\n",
      "loss: 0.000001  [ 2560/60000]\n",
      "loss: 0.000000  [ 3200/60000]\n",
      "loss: 0.000000  [ 3840/60000]\n",
      "loss: 0.000000  [ 4480/60000]\n",
      "loss: 0.000000  [ 5120/60000]\n",
      "loss: 0.000000  [ 5760/60000]\n",
      "loss: 7.194639  [ 6400/60000]\n",
      "loss: 5.615434  [ 7040/60000]\n",
      "loss: 3.928238  [ 7680/60000]\n",
      "loss: 2.563159  [ 8320/60000]\n",
      "loss: 1.594102  [ 8960/60000]\n",
      "loss: 0.983429  [ 9600/60000]\n",
      "loss: 0.621121  [10240/60000]\n",
      "loss: 0.409284  [10880/60000]\n",
      "loss: 0.286960  [11520/60000]\n",
      "loss: 0.214131  [12160/60000]\n",
      "loss: 3.527124  [12800/60000]\n",
      "loss: 1.457805  [13440/60000]\n",
      "loss: 0.468894  [14080/60000]\n",
      "loss: 0.127355  [14720/60000]\n",
      "loss: 0.018079  [15360/60000]\n",
      "loss: 0.000795  [16000/60000]\n",
      "loss: 0.000220  [16640/60000]\n",
      "loss: 0.009381  [17280/60000]\n",
      "loss: 0.047201  [17920/60000]\n",
      "loss: 0.076694  [18560/60000]\n",
      "loss: 3.479500  [19200/60000]\n",
      "loss: 2.284616  [19840/60000]\n",
      "loss: 1.750637  [20480/60000]\n",
      "loss: 1.153878  [21120/60000]\n",
      "loss: 0.695157  [21760/60000]\n",
      "loss: 0.461287  [22400/60000]\n",
      "loss: 0.342143  [23040/60000]\n",
      "loss: 0.234501  [23680/60000]\n",
      "loss: 0.217011  [24320/60000]\n",
      "loss: 4.712353  [24960/60000]\n",
      "loss: 1.678658  [25600/60000]\n",
      "loss: 0.136111  [26240/60000]\n",
      "loss: 0.072891  [26880/60000]\n",
      "loss: 0.000381  [27520/60000]\n",
      "loss: 0.001552  [28160/60000]\n",
      "loss: 0.016227  [28800/60000]\n",
      "loss: 0.000565  [29440/60000]\n",
      "loss: 0.019267  [30080/60000]\n",
      "loss: 5.487148  [30720/60000]\n",
      "loss: 2.088026  [31360/60000]\n",
      "loss: 1.146757  [32000/60000]\n",
      "loss: 0.160710  [32640/60000]\n",
      "loss: 0.002665  [33280/60000]\n",
      "loss: 0.000822  [33920/60000]\n",
      "loss: 0.001436  [34560/60000]\n",
      "loss: 0.000379  [35200/60000]\n",
      "loss: 0.000400  [35840/60000]\n",
      "loss: 0.911210  [36480/60000]\n",
      "loss: 0.009380  [37120/60000]\n",
      "loss: 0.000395  [37760/60000]\n",
      "loss: 0.000017  [38400/60000]\n",
      "loss: 0.002049  [39040/60000]\n",
      "loss: 0.000018  [39680/60000]\n",
      "loss: 0.000010  [40320/60000]\n",
      "loss: 0.000114  [40960/60000]\n",
      "loss: 0.000012  [41600/60000]\n",
      "loss: 3.743647  [42240/60000]\n",
      "loss: 2.521879  [42880/60000]\n",
      "loss: 0.169121  [43520/60000]\n",
      "loss: 0.004181  [44160/60000]\n",
      "loss: 0.003360  [44800/60000]\n",
      "loss: 0.000107  [45440/60000]\n",
      "loss: 0.000138  [46080/60000]\n",
      "loss: 0.000190  [46720/60000]\n",
      "loss: 0.000137  [47360/60000]\n",
      "loss: 0.000038  [48000/60000]\n",
      "loss: 1.946815  [48640/60000]\n",
      "loss: 1.347068  [49280/60000]\n",
      "loss: 0.146495  [49920/60000]\n",
      "loss: 0.181311  [50560/60000]\n",
      "loss: 0.051809  [51200/60000]\n",
      "loss: 0.000059  [51840/60000]\n",
      "loss: 0.001100  [52480/60000]\n",
      "loss: 0.000457  [53120/60000]\n",
      "loss: 0.000210  [53760/60000]\n",
      "loss: 15.170379  [54400/60000]\n",
      "loss: 1.696596  [55040/60000]\n",
      "loss: 1.091440  [55680/60000]\n",
      "loss: 0.604186  [56320/60000]\n",
      "loss: 0.391359  [56960/60000]\n",
      "loss: 0.251258  [57600/60000]\n",
      "loss: 0.194416  [58240/60000]\n",
      "loss: 0.163378  [58880/60000]\n",
      "loss: 0.141236  [59520/60000]\n",
      "Training done!\n"
     ]
    }
   ],
   "source": [
    "def train(dataloader, model, loss_fn, optimizer):\n",
    "    model.train()\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        X, y = X.to(device), y.to(device)\n",
    "\n",
    "        # Compute prediction error\n",
    "        pred = model(X)\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        # Backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        if batch % 10 == 0:\n",
    "            loss, current = loss.item(), batch * len(X)\n",
    "            print(f\"loss: {loss:>7f}  [{current:>5d}/{len(dataloader.dataset):>5d}]\")\n",
    "\n",
    "# Training the model\n",
    "epochs = 5\n",
    "for epoch in range(epochs):\n",
    "    print(f\"Epoch {epoch+1}\\n-------------------------------\")\n",
    "    train(train_loader, model, loss_fn, optimizer)\n",
    "print(\"Training done!\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate(dataloader, model, loss_fn):\n",
    "    model.eval()\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    loss, correct = 0, 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            pred = model(X)\n",
    "            loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "\n",
    "    loss /= num_batches\n",
    "    correct /= size\n",
    "    accuracy = 100 * correct\n",
    "    print(f\"Test Error: \\n Accuracy: {(accuracy):>0.1f}%, Avg loss: {loss:>8f} \\n\")\n",
    "    return accuracy\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Error: \n",
      " Accuracy: 10.1%, Avg loss: 5.157917 \n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "10.11"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validate(test_loader, model, loss_fn)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}