L14/2-resnet-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "STAT 453: Deep Learning (Spring 2021)  \n",
    "Instructor: Sebastian Raschka (sraschka@wisc.edu)  \n",
    "\n",
    "Course website: http://pages.stat.wisc.edu/~sraschka/teaching/stat453-ss2021/  \n",
    "GitHub repository: https://github.com/rasbt/stat453-deep-learning-ss21\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sebastian Raschka\n",
      "\n",
      "Python implementation: CPython\n",
      "Python version       : 3.9.2\n",
      "IPython version      : 7.21.0\n",
      "\n",
      "torch: 1.9.0a0+d819a21\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%load_ext watermark\n",
    "%watermark -a 'Sebastian Raschka' -v -p torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Runs on CPU or GPU (if available)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Convolutional ResNet and Residual Blocks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please note that this example does not implement a really deep ResNet as described in literature but rather illustrates how the residual blocks described in He et al. [1] can be implemented in PyTorch.\n",
    "\n",
    "- [1] He, Kaiming, et al. \"Deep residual learning for image recognition.\" *Proceedings of the IEEE conference on computer vision and pattern recognition*. 2016."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings and Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image batch dimensions: torch.Size([128, 1, 28, 28])\n",
      "Image label dimensions: torch.Size([128])\n"
     ]
    }
   ],
   "source": [
    "##########################\n",
    "### SETTINGS\n",
    "##########################\n",
    "\n",
    "# Device\n",
    "device = torch.device(\"cuda:1\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Hyperparameters\n",
    "random_seed = 123\n",
    "learning_rate = 0.01\n",
    "num_epochs = 10\n",
    "batch_size = 128\n",
    "\n",
    "# Architecture\n",
    "num_classes = 10\n",
    "\n",
    "\n",
    "##########################\n",
    "### MNIST DATASET\n",
    "##########################\n",
    "\n",
    "# Note transforms.ToTensor() scales input images\n",
    "# to 0-1 range\n",
    "train_dataset = datasets.MNIST(root='data', \n",
    "                               train=True, \n",
    "                               transform=transforms.ToTensor(),\n",
    "                               download=True)\n",
    "\n",
    "test_dataset = datasets.MNIST(root='data', \n",
    "                              train=False, \n",
    "                              transform=transforms.ToTensor())\n",
    "\n",
    "\n",
    "train_loader = DataLoader(dataset=train_dataset, \n",
    "                          batch_size=batch_size, \n",
    "                          shuffle=True)\n",
    "\n",
    "test_loader = DataLoader(dataset=test_dataset, \n",
    "                         batch_size=batch_size, \n",
    "                         shuffle=False)\n",
    "\n",
    "# Checking the dataset\n",
    "for images, labels in train_loader:  \n",
    "    print('Image batch dimensions:', images.shape)\n",
    "    print('Image label dimensions:', labels.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ResNet with identity blocks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code implements the residual blocks with skip connections such that the input passed via the shortcut matches the dimensions of the main path's output, which allows the network to learn identity functions. Such a residual block is illustrated below:\n",
    "\n",
    "![](./2-resnet-ex/resnet-ex-1-1.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "##########################\n",
    "### MODEL\n",
    "##########################\n",
    "\n",
    "\n",
    "class ConvNet(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, num_classes):\n",
    "        super(ConvNet, self).__init__()\n",
    "        \n",
    "        #########################\n",
    "        ### 1st residual block\n",
    "        #########################\n",
    "        \n",
    "        self.block_1 = torch.nn.Sequential(\n",
    "                torch.nn.Conv2d(in_channels=1,\n",
    "                                out_channels=4,\n",
    "                                kernel_size=(1, 1),\n",
    "                                stride=(1, 1),\n",
    "                                padding=0),\n",
    "                torch.nn.BatchNorm2d(4),\n",
    "                torch.nn.ReLU(inplace=True),\n",
    "                torch.nn.Conv2d(in_channels=4,\n",
    "                                out_channels=1,\n",
    "                                kernel_size=(3, 3),\n",
    "                                stride=(1, 1),\n",
    "                                padding=1),\n",
    "                torch.nn.BatchNorm2d(1)\n",
    "        )\n",
    "        \n",
    "        self.block_2 = torch.nn.Sequential(\n",
    "                torch.nn.Conv2d(in_channels=1,\n",
    "                                out_channels=4,\n",
    "                                kernel_size=(1, 1),\n",
    "                                stride=(1, 1),\n",
    "                                padding=0),\n",
    "                torch.nn.BatchNorm2d(4),\n",
    "                torch.nn.ReLU(inplace=True),\n",
    "                torch.nn.Conv2d(in_channels=4,\n",
    "                                out_channels=1,\n",
    "                                kernel_size=(3, 3),\n",
    "                                stride=(1, 1),\n",
    "                                padding=1),\n",
    "                torch.nn.BatchNorm2d(1)\n",
    "        )\n",
    "\n",
    "        #########################\n",
    "        ### Fully connected\n",
    "        #########################        \n",
    "        self.linear_1 = torch.nn.Linear(1*28*28, num_classes)\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        #########################\n",
    "        ### 1st residual block\n",
    "        #########################\n",
    "        shortcut = x\n",
    "        x = self.block_1(x)\n",
    "        x = torch.nn.functional.relu(x + shortcut)\n",
    "        \n",
    "        #########################\n",
    "        ### 2nd residual block\n",
    "        #########################\n",
    "        shortcut = x\n",
    "        x = self.block_2(x)\n",
    "        x = torch.nn.functional.relu(x + shortcut)\n",
    "        \n",
    "        #########################\n",
    "        ### Fully connected\n",
    "        #########################\n",
    "        logits = self.linear_1(x.view(-1,  1*28*28))\n",
    "        return logits\n",
    "\n",
    "    \n",
    "torch.manual_seed(random_seed)\n",
    "model = ConvNet(num_classes=num_classes)\n",
    "model = model.to(device)\n",
    "    \n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001/010 | Batch 000/469 | Cost: 2.6938\n",
      "Epoch: 001/010 | Batch 250/469 | Cost: 0.3444\n",
      "Epoch: 001/010 training accuracy: 91.21%\n",
      "Time elapsed: 0.29 min\n",
      "Epoch: 002/010 | Batch 000/469 | Cost: 0.4317\n",
      "Epoch: 002/010 | Batch 250/469 | Cost: 0.4750\n",
      "Epoch: 002/010 training accuracy: 92.69%\n",
      "Time elapsed: 0.57 min\n",
      "Epoch: 003/010 | Batch 000/469 | Cost: 0.2430\n",
      "Epoch: 003/010 | Batch 250/469 | Cost: 0.2722\n",
      "Epoch: 003/010 training accuracy: 92.56%\n",
      "Time elapsed: 0.86 min\n",
      "Epoch: 004/010 | Batch 000/469 | Cost: 0.2547\n",
      "Epoch: 004/010 | Batch 250/469 | Cost: 0.2330\n",
      "Epoch: 004/010 training accuracy: 92.79%\n",
      "Time elapsed: 1.15 min\n",
      "Epoch: 005/010 | Batch 000/469 | Cost: 0.0948\n",
      "Epoch: 005/010 | Batch 250/469 | Cost: 0.2342\n",
      "Epoch: 005/010 training accuracy: 93.24%\n",
      "Time elapsed: 1.43 min\n",
      "Epoch: 006/010 | Batch 000/469 | Cost: 0.2508\n",
      "Epoch: 006/010 | Batch 250/469 | Cost: 0.2552\n",
      "Epoch: 006/010 training accuracy: 93.11%\n",
      "Time elapsed: 1.72 min\n",
      "Epoch: 007/010 | Batch 000/469 | Cost: 0.0941\n",
      "Epoch: 007/010 | Batch 250/469 | Cost: 0.3898\n",
      "Epoch: 007/010 training accuracy: 93.14%\n",
      "Time elapsed: 2.00 min\n",
      "Epoch: 008/010 | Batch 000/469 | Cost: 0.2865\n",
      "Epoch: 008/010 | Batch 250/469 | Cost: 0.1691\n",
      "Epoch: 008/010 training accuracy: 93.73%\n",
      "Time elapsed: 2.29 min\n",
      "Epoch: 009/010 | Batch 000/469 | Cost: 0.3685\n",
      "Epoch: 009/010 | Batch 250/469 | Cost: 0.1544\n",
      "Epoch: 009/010 training accuracy: 92.88%\n",
      "Time elapsed: 2.58 min\n",
      "Epoch: 010/010 | Batch 000/469 | Cost: 0.3586\n",
      "Epoch: 010/010 | Batch 250/469 | Cost: 0.4787\n",
      "Epoch: 010/010 training accuracy: 93.65%\n",
      "Time elapsed: 2.87 min\n",
      "Total Training Time: 2.87 min\n"
     ]
    }
   ],
   "source": [
    "def compute_accuracy(model, data_loader):\n",
    "    correct_pred, num_examples = 0, 0\n",
    "    for i, (features, targets) in enumerate(data_loader):            \n",
    "        features = features.to(device)\n",
    "        targets = targets.to(device)\n",
    "        logits = model(features)\n",
    "        _, predicted_labels = torch.max(logits, 1)\n",
    "        num_examples += targets.size(0)\n",
    "        correct_pred += (predicted_labels == targets).sum()\n",
    "    return correct_pred.float()/num_examples * 100\n",
    "\n",
    "\n",
    "start_time = time.time()\n",
    "for epoch in range(num_epochs):\n",
    "    model = model.train()\n",
    "    for batch_idx, (features, targets) in enumerate(train_loader):\n",
    "        \n",
    "        features = features.to(device)\n",
    "        targets = targets.to(device)\n",
    "        \n",
    "        ### FORWARD AND BACK PROP\n",
    "        logits = model(features)\n",
    "        cost = torch.nn.functional.cross_entropy(logits, targets)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        cost.backward()\n",
    "        \n",
    "        ### UPDATE MODEL PARAMETERS\n",
    "        optimizer.step()\n",
    "        \n",
    "        ### LOGGING\n",
    "        if not batch_idx % 250:\n",
    "            print ('Epoch: %03d/%03d | Batch %03d/%03d | Cost: %.4f' \n",
    "                   %(epoch+1, num_epochs, batch_idx, \n",
    "                     len(train_loader), cost))\n",
    "\n",
    "    model = model.eval() # eval mode to prevent upd. batchnorm params during inference\n",
    "    with torch.set_grad_enabled(False): # save memory during inference\n",
    "        print('Epoch: %03d/%03d training accuracy: %.2f%%' % (\n",
    "              epoch+1, num_epochs, \n",
    "              compute_accuracy(model, train_loader)))\n",
    "\n",
    "    print('Time elapsed: %.2f min' % ((time.time() - start_time)/60))\n",
    "    \n",
    "print('Total Training Time: %.2f min' % ((time.time() - start_time)/60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy: 92.16%\n"
     ]
    }
   ],
   "source": [
    "print('Test accuracy: %.2f%%' % (compute_accuracy(model, test_loader)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ResNet with convolutional blocks for resizing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code implements the residual blocks with skip connections such that the input passed via the shortcut matches is resized to dimensions of the main path's output. Such a residual block is illustrated below:\n",
    "\n",
    "![](./2-resnet-ex/resnet-ex-1-2.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResidualBlock(torch.nn.Module):\n",
    "    \"\"\" Helper Class\"\"\"\n",
    "\n",
    "    def __init__(self, channels):\n",
    "        \n",
    "        super(ResidualBlock, self).__init__()\n",
    "        \n",
    "        self.block = torch.nn.Sequential(\n",
    "                torch.nn.Conv2d(in_channels=channels[0],\n",
    "                                out_channels=channels[1],\n",
    "                                kernel_size=(3, 3),\n",
    "                                stride=(2, 2),\n",
    "                                padding=1),\n",
    "                torch.nn.BatchNorm2d(channels[1]),\n",
    "                torch.nn.ReLU(inplace=True),\n",
    "                torch.nn.Conv2d(in_channels=channels[1],\n",
    "                                out_channels=channels[2],\n",
    "                                kernel_size=(1, 1),\n",
    "                                stride=(1, 1),\n",
    "                                padding=0),   \n",
    "                torch.nn.BatchNorm2d(channels[2])\n",
    "        )\n",
    "\n",
    "        self.shortcut = torch.nn.Sequential(\n",
    "                torch.nn.Conv2d(in_channels=channels[0],\n",
    "                                out_channels=channels[2],\n",
    "                                kernel_size=(1, 1),\n",
    "                                stride=(2, 2),\n",
    "                                padding=0),\n",
    "                torch.nn.BatchNorm2d(channels[2])\n",
    "        )\n",
    "            \n",
    "    def forward(self, x):\n",
    "        shortcut = x\n",
    "        \n",
    "        block = self.block(x)\n",
    "        shortcut = self.shortcut(x)    \n",
    "        x = torch.nn.functional.relu(block+shortcut)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "##########################\n",
    "### MODEL\n",
    "##########################\n",
    "\n",
    "\n",
    "\n",
    "class ConvNet(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, num_classes):\n",
    "        super(ConvNet, self).__init__()\n",
    "        \n",
    "        self.residual_block_1 = ResidualBlock(channels=[1, 4, 8])\n",
    "        self.residual_block_2 = ResidualBlock(channels=[8, 16, 32])\n",
    "    \n",
    "        self.linear_1 = torch.nn.Linear(7*7*32, num_classes)\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "\n",
    "        out = self.residual_block_1(x)\n",
    "        out = self.residual_block_2(out)\n",
    "         \n",
    "        logits = self.linear_1(out.view(-1, 7*7*32))\n",
    "        return logits\n",
    "\n",
    "    \n",
    "torch.manual_seed(random_seed)\n",
    "model = ConvNet(num_classes=num_classes)\n",
    "\n",
    "model.to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001/010 | Batch 000/468 | Cost: 2.4548\n",
      "Epoch: 001/010 | Batch 050/468 | Cost: 0.2955\n",
      "Epoch: 001/010 | Batch 100/468 | Cost: 0.1596\n",
      "Epoch: 001/010 | Batch 150/468 | Cost: 0.1105\n",
      "Epoch: 001/010 | Batch 200/468 | Cost: 0.1173\n",
      "Epoch: 001/010 | Batch 250/468 | Cost: 0.1747\n",
      "Epoch: 001/010 | Batch 300/468 | Cost: 0.0986\n",
      "Epoch: 001/010 | Batch 350/468 | Cost: 0.1371\n",
      "Epoch: 001/010 | Batch 400/468 | Cost: 0.1583\n",
      "Epoch: 001/010 | Batch 450/468 | Cost: 0.1172\n",
      "Epoch: 001/010 training accuracy: 97.29%\n",
      "Epoch: 002/010 | Batch 000/468 | Cost: 0.1137\n",
      "Epoch: 002/010 | Batch 050/468 | Cost: 0.0812\n",
      "Epoch: 002/010 | Batch 100/468 | Cost: 0.0268\n",
      "Epoch: 002/010 | Batch 150/468 | Cost: 0.0891\n",
      "Epoch: 002/010 | Batch 200/468 | Cost: 0.0918\n",
      "Epoch: 002/010 | Batch 250/468 | Cost: 0.0894\n",
      "Epoch: 002/010 | Batch 300/468 | Cost: 0.0487\n",
      "Epoch: 002/010 | Batch 350/468 | Cost: 0.0395\n",
      "Epoch: 002/010 | Batch 400/468 | Cost: 0.0729\n",
      "Epoch: 002/010 | Batch 450/468 | Cost: 0.0465\n",
      "Epoch: 002/010 training accuracy: 97.92%\n",
      "Epoch: 003/010 | Batch 000/468 | Cost: 0.0624\n",
      "Epoch: 003/010 | Batch 050/468 | Cost: 0.0483\n",
      "Epoch: 003/010 | Batch 100/468 | Cost: 0.0574\n",
      "Epoch: 003/010 | Batch 150/468 | Cost: 0.0765\n",
      "Epoch: 003/010 | Batch 200/468 | Cost: 0.1354\n",
      "Epoch: 003/010 | Batch 250/468 | Cost: 0.1235\n",
      "Epoch: 003/010 | Batch 300/468 | Cost: 0.1017\n",
      "Epoch: 003/010 | Batch 350/468 | Cost: 0.0681\n",
      "Epoch: 003/010 | Batch 400/468 | Cost: 0.0438\n",
      "Epoch: 003/010 | Batch 450/468 | Cost: 0.1454\n",
      "Epoch: 003/010 training accuracy: 98.13%\n",
      "Epoch: 004/010 | Batch 000/468 | Cost: 0.0770\n",
      "Epoch: 004/010 | Batch 050/468 | Cost: 0.0204\n",
      "Epoch: 004/010 | Batch 100/468 | Cost: 0.0570\n",
      "Epoch: 004/010 | Batch 150/468 | Cost: 0.0614\n",
      "Epoch: 004/010 | Batch 200/468 | Cost: 0.0572\n",
      "Epoch: 004/010 | Batch 250/468 | Cost: 0.0353\n",
      "Epoch: 004/010 | Batch 300/468 | Cost: 0.0270\n",
      "Epoch: 004/010 | Batch 350/468 | Cost: 0.0930\n",
      "Epoch: 004/010 | Batch 400/468 | Cost: 0.0832\n",
      "Epoch: 004/010 | Batch 450/468 | Cost: 0.0512\n",
      "Epoch: 004/010 training accuracy: 98.57%\n",
      "Epoch: 005/010 | Batch 000/468 | Cost: 0.0429\n",
      "Epoch: 005/010 | Batch 050/468 | Cost: 0.0750\n",
      "Epoch: 005/010 | Batch 100/468 | Cost: 0.0065\n",
      "Epoch: 005/010 | Batch 150/468 | Cost: 0.0870\n",
      "Epoch: 005/010 | Batch 200/468 | Cost: 0.0320\n",
      "Epoch: 005/010 | Batch 250/468 | Cost: 0.0463\n",
      "Epoch: 005/010 | Batch 300/468 | Cost: 0.0611\n",
      "Epoch: 005/010 | Batch 350/468 | Cost: 0.0231\n",
      "Epoch: 005/010 | Batch 400/468 | Cost: 0.0404\n",
      "Epoch: 005/010 | Batch 450/468 | Cost: 0.0195\n",
      "Epoch: 005/010 training accuracy: 98.85%\n",
      "Epoch: 006/010 | Batch 000/468 | Cost: 0.0530\n",
      "Epoch: 006/010 | Batch 050/468 | Cost: 0.0614\n",
      "Epoch: 006/010 | Batch 100/468 | Cost: 0.0070\n",
      "Epoch: 006/010 | Batch 150/468 | Cost: 0.0147\n",
      "Epoch: 006/010 | Batch 200/468 | Cost: 0.1086\n",
      "Epoch: 006/010 | Batch 250/468 | Cost: 0.0107\n",
      "Epoch: 006/010 | Batch 300/468 | Cost: 0.0102\n",
      "Epoch: 006/010 | Batch 350/468 | Cost: 0.0190\n",
      "Epoch: 006/010 | Batch 400/468 | Cost: 0.1027\n",
      "Epoch: 006/010 | Batch 450/468 | Cost: 0.0590\n",
      "Epoch: 006/010 training accuracy: 99.03%\n",
      "Epoch: 007/010 | Batch 000/468 | Cost: 0.0061\n",
      "Epoch: 007/010 | Batch 050/468 | Cost: 0.0112\n",
      "Epoch: 007/010 | Batch 100/468 | Cost: 0.0361\n",
      "Epoch: 007/010 | Batch 150/468 | Cost: 0.0968\n",
      "Epoch: 007/010 | Batch 200/468 | Cost: 0.0246\n",
      "Epoch: 007/010 | Batch 250/468 | Cost: 0.0144\n",
      "Epoch: 007/010 | Batch 300/468 | Cost: 0.0238\n",
      "Epoch: 007/010 | Batch 350/468 | Cost: 0.0086\n",
      "Epoch: 007/010 | Batch 400/468 | Cost: 0.0220\n",
      "Epoch: 007/010 | Batch 450/468 | Cost: 0.0672\n",
      "Epoch: 007/010 training accuracy: 99.04%\n",
      "Epoch: 008/010 | Batch 000/468 | Cost: 0.0551\n",
      "Epoch: 008/010 | Batch 050/468 | Cost: 0.0156\n",
      "Epoch: 008/010 | Batch 100/468 | Cost: 0.0025\n",
      "Epoch: 008/010 | Batch 150/468 | Cost: 0.0116\n",
      "Epoch: 008/010 | Batch 200/468 | Cost: 0.0093\n",
      "Epoch: 008/010 | Batch 250/468 | Cost: 0.0275\n",
      "Epoch: 008/010 | Batch 300/468 | Cost: 0.0581\n",
      "Epoch: 008/010 | Batch 350/468 | Cost: 0.0211\n",
      "Epoch: 008/010 | Batch 400/468 | Cost: 0.0224\n",
      "Epoch: 008/010 | Batch 450/468 | Cost: 0.0520\n",
      "Epoch: 008/010 training accuracy: 99.19%\n",
      "Epoch: 009/010 | Batch 000/468 | Cost: 0.0218\n",
      "Epoch: 009/010 | Batch 050/468 | Cost: 0.0742\n",
      "Epoch: 009/010 | Batch 100/468 | Cost: 0.0200\n",
      "Epoch: 009/010 | Batch 150/468 | Cost: 0.1219\n",
      "Epoch: 009/010 | Batch 200/468 | Cost: 0.0438\n",
      "Epoch: 009/010 | Batch 250/468 | Cost: 0.0195\n",
      "Epoch: 009/010 | Batch 300/468 | Cost: 0.0356\n",
      "Epoch: 009/010 | Batch 350/468 | Cost: 0.0959\n",
      "Epoch: 009/010 | Batch 400/468 | Cost: 0.0272\n",
      "Epoch: 009/010 | Batch 450/468 | Cost: 0.1140\n",
      "Epoch: 009/010 training accuracy: 99.09%\n",
      "Epoch: 010/010 | Batch 000/468 | Cost: 0.0308\n",
      "Epoch: 010/010 | Batch 050/468 | Cost: 0.0333\n",
      "Epoch: 010/010 | Batch 100/468 | Cost: 0.0286\n",
      "Epoch: 010/010 | Batch 150/468 | Cost: 0.0202\n",
      "Epoch: 010/010 | Batch 200/468 | Cost: 0.0193\n",
      "Epoch: 010/010 | Batch 250/468 | Cost: 0.0057\n",
      "Epoch: 010/010 | Batch 300/468 | Cost: 0.0331\n",
      "Epoch: 010/010 | Batch 350/468 | Cost: 0.0200\n",
      "Epoch: 010/010 | Batch 400/468 | Cost: 0.0588\n",
      "Epoch: 010/010 | Batch 450/468 | Cost: 0.0202\n",
      "Epoch: 010/010 training accuracy: 99.07%\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    model = model.train()\n",
    "    for batch_idx, (features, targets) in enumerate(train_loader):\n",
    "        \n",
    "        features = features.to(device)\n",
    "        targets = targets.to(device)\n",
    "            \n",
    "        ### FORWARD AND BACK PROP\n",
    "        logits = model(features)\n",
    "        cost = torch.nn.functional.cross_entropy(logits, targets)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        cost.backward()\n",
    "        \n",
    "        ### UPDATE MODEL PARAMETERS\n",
    "        optimizer.step()\n",
    "        \n",
    "        ### LOGGING\n",
    "        if not batch_idx % 50:\n",
    "            print ('Epoch: %03d/%03d | Batch %03d/%03d | Cost: %.4f' \n",
    "                   %(epoch+1, num_epochs, batch_idx, \n",
    "                     len(train_dataset)//batch_size, cost))\n",
    "\n",
    "    model = model.eval() # eval mode to prevent upd. batchnorm params during inference\n",
    "    with torch.set_grad_enabled(False): # save memory during inference\n",
    "        print('Epoch: %03d/%03d training accuracy: %.2f%%' % (\n",
    "              epoch+1, num_epochs, \n",
    "              compute_accuracy(model, train_loader)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy: 98.13%\n"
     ]
    }
   ],
   "source": [
    "print('Test accuracy: %.2f%%' % (compute_accuracy(model, test_loader)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.2"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}