2. MNIST Resnet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3952c183",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.io import read_image\n",
    "from torchvision.transforms import ToTensor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6770f008",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c4042bd7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda \n"
     ]
    }
   ],
   "source": [
    "# Get cpu or gpu device for training.\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using device: {device} \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fbba117d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Current Cuda Device: 0\n",
      "Device 0 name: NVIDIA GeForce GTX 1060\n"
     ]
    }
   ],
   "source": [
    "print(\"Current Cuda Device: {}\".format(torch.cuda.current_device()))\n",
    "print(\"Device 0 name: {}\".format(torch.cuda.get_device_name(0)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "af837ff9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Default PyTorch dataset framework; img_dir must be specified and csv rows must be formatted as: filename,label\n",
    "class CustomImageDataset(Dataset):\n",
    "    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):\n",
    "        self.img_labels = pd.read_csv(annotations_file)\n",
    "        self.img_dir = img_dir\n",
    "        self.transform = transform\n",
    "        self.target_transform = target_transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_labels)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])\n",
    "        image = read_image(img_path)\n",
    "        label = self.img_labels.iloc[idx, 1]\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        if self.target_transform:\n",
    "            label = self.target_transform(label)\n",
    "        return image, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "92c9f184",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define general NN model\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, 1)\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.dropout = nn.Dropout(0.25)\n",
    "        self.linear = nn.Linear(21632, 500)\n",
    "        self.linear2 = nn.Linear(500, 10)\n",
    "        self.leakyrelu = nn.LeakyReLU(negative_slope=0.01)\n",
    "        self.relu = nn.ReLU()\n",
    "#         self.stack = nn.Sequential(\n",
    "#             nn.ReLU(),\n",
    "#             nn.MaxPool2d(2)\n",
    "#         )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.leakyrelu(x)\n",
    "        x = self.dropout(x)\n",
    "        x = self.flatten(x)\n",
    "        x = self.linear(x)\n",
    "        x = self.relu(x)\n",
    "        logits = self.linear2(x)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "412dd760",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = NeuralNetwork()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "058dbcd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Per PyTorch: All pre-trained models expect input images normalized in the same way:\n",
    "# i.e. mini-batches of 3-channel RGB images of shape (3 x H x W), where H and W are expected to be at least 224\n",
    "# The images have to be loaded in to a range of [0, 1] and then normalized using\n",
    "# mean = [0.485, 0.456, 0.406] and std = [0.229, 0.224, 0.225].\n",
    "preprocess = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "7774c134",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in C:\\Users\\thoma/.cache\\torch\\hub\\pytorch_vision_v0.10.0\n"
     ]
    }
   ],
   "source": [
    "# If using residual net pretrained on ImageNet data\n",
    "model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet152', pretrained=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "3df85d70",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ResNet(\n",
      "  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
      "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (relu): ReLU(inplace=True)\n",
      "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "  (layer1): Sequential(\n",
      "    (0): Bottleneck(\n",
      "      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (1): Bottleneck(\n",
      "      (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (2): Bottleneck(\n",
      "      (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "  )\n",
      "  (layer2): Sequential(\n",
      "    (0): Bottleneck(\n",
      "      (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (1): Bottleneck(\n",
      "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (2): Bottleneck(\n",
      "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (3): Bottleneck(\n",
      "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (4): Bottleneck(\n",
      "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (5): Bottleneck(\n",
      "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (6): Bottleneck(\n",
      "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (7): Bottleneck(\n",
      "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "  )\n",
      "  (layer3): Sequential(\n",
      "    (0): Bottleneck(\n",
      "      (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (1): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (2): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (3): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (4): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (5): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (6): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (7): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (8): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (9): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (10): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (11): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (12): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (13): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (14): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (15): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (16): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (17): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (18): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (19): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (20): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (21): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (22): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (23): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (24): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (25): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (26): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (27): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (28): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (29): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (30): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (31): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (32): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (33): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (34): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (35): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "  )\n",
      "  (layer4): Sequential(\n",
      "    (0): Bottleneck(\n",
      "      (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (1): Bottleneck(\n",
      "      (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "    (2): Bottleneck(\n",
      "      (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "    )\n",
      "  )\n",
      "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
      "  (fc): Linear(in_features=2048, out_features=1000, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model.to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3f7c7101",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr = 0.005)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d31f01ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(dataloader, model, loss_fn, optimizer):\n",
    "    size = len(dataloader.dataset)\n",
    "    model.train()\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        \n",
    "        # Preprocessing needed for pretrained models like ResNet\n",
    "        input_tensor = preprocess(input_image)\n",
    "        input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model\n",
    "        \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 % 100 == 0:\n",
    "            loss, current = loss.item(), batch * len(X)\n",
    "            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "65fd8b8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(dataloader, model, loss_fn):\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    model.eval()\n",
    "    test_loss, correct = 0, 0\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",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "    test_loss /= num_batches\n",
    "    correct /= size\n",
    "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40c8dd24",
   "metadata": {},
   "source": [
    "### Testing on Public MNIST Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "0369e574",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_data = datasets.MNIST(\n",
    "    root=\"data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "6b3f0624",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = datasets.MNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "53dcd214",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of X [N, C, H, W]: torch.Size([64, 1, 28, 28])\n",
      "Shape of y: torch.Size([64]) torch.int64\n"
     ]
    }
   ],
   "source": [
    "batch_size = 64\n",
    "\n",
    "# Create data loaders.\n",
    "train_dataloader = DataLoader(training_data, batch_size=batch_size)\n",
    "test_dataloader = DataLoader(test_data, batch_size=batch_size)\n",
    "\n",
    "for X, y in train_dataloader:\n",
    "    print(f\"Shape of X [N, C, H, W]: {X.shape}\")\n",
    "    print(f\"Shape of y: {y.shape} {y.dtype}\")\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "4ef1bb72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAS4AAAD8CAYAAADJwUnTAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAVoElEQVR4nO3df5BdZX3H8fcnS0IgYCHEQEyCIIa22KnBLsERsbH4IzC2gVEpsSoy2kiHdMqMf0CZttBxOsUfiFiRzAopMKNQKyipEw0/xoqOggkYISGgMUayJE2IQflNkt1v/zg3evfeveee3b17z3l2P6+ZM3vPec557neu5svzPOc5z1FEYGaWkillB2BmNlJOXGaWHCcuM0uOE5eZJceJy8yS48RlZslx4jKzcSNplaTdkja2KJekL0jaIukRSW8qUq8Tl5mNp5uBJTnlZwMLatty4IYilTpxmdm4iYj7gb05pywFbo3MA8BRkua0q/eQTgVYxDQdGtOZ0c2vNJtUXuYF9sUrGksd7377jPj13oFC5z70yCubgJfrDvVFRN8Ivm4usL1uv792bGfeRWNKXJKWANcBPcCNEXF13vnTmcHpOmssX2lmOR6M+8Zcx569Azy4dl6hc6fO+cXLEdE7hq8bLsm2fQ5x1IlLUg9wPfBOsiy5TtLqiHhstHWaWRUEAzHYrS/rB+bX7c8DdrS7aCxjXIuALRGxNSL2AbeT9VfNLGEBDBKFtg5YDXy4dnfxzcBvIyK3mwhj6yoO1zc9vfEkScvJ7hYwncPH8HVm1i2DdKbFJek2YDEwS1I/cCUwFSAiVgJrgHOALcCLwEVF6h1L4irUN60N1PUBvEozvYaOWcUFwf4OdRUjYlmb8gAuGWm9Y0lco+qbmlm1BTDQmW7guBnLGNc6YIGkEyVNAy4g66+aWeK6OMY1KqNucUXEAUkrgLVk0yFWRcSmjkVmZqUIYKDiKyOPaR5XRKwhG1wzswmka5MhRqmrM+fNrPqCqPwYlxOXmQ0RAfurnbecuMyskRgYdrZTdThxmdkQAQy6xWVmqXGLy8ySkk1AdeIys4QEsD+qvcaoE5eZDRGIgYovjuzEZWZNBsNdRTNLiMe4zCxBYsBjXGaWkmwFVCcuM0tIhNgXPWWHkcuJy8yaDHqMy8xSkg3Ou6toZknx4LyZJcaD82aWpAFPQDWzlARif1Q7NVQ7OjPrOg/Om1lyArmraGbp8eC8mSUlAk+HMLO0ZIPzfuTHzBLjwXkzS0ogLyRoZulxi8vMkpK9V9GJy8yS4jdZm5Xmhfed3rLsU5++IffaT57/4dzyWL9xVDGlIHs92QS+qyhpG/AcMAAciIjeTgRlZuWJUOW7ip2I7u0RsdBJy2ziGIgphbYiJC2R9ISkLZIuH6b8DyT9j6SfStok6aJ2dVY7rZpZ12XrcanQ1o6kHuB64GzgFGCZpFMaTrsEeCwi3ggsBq6RNC2v3rEmrgDulvSQpOUtAl8uab2k9ft5ZYxfZ2bjT51scS0CtkTE1ojYB9wOLG04J4AjJQk4AtgLHMirdKyD82dExA5Js4F7JD0eEfcPiSiiD+gDeJVmxhi/z8zGWTYdovBdxVmS1tft99X+zR80F9het98PNN41+SKwGtgBHAn8dUQM5n3pmBJXROyo/d0t6Rtk2fX+/KvMrMpG+Kzinjbj28NlwMYGzLuBDcBfACeRNYK+HxHPtqp01F1FSTMkHXnwM/AuYOLeIzabRAaZUmgroB+YX7c/j6xlVe8i4M7IbAF+CfxRXqVjaXEdC3wj65ZyCPDViPjOGOobVy8tXZRffkz+f2FmrvpRJ8OxLtjd2/of1ie3/WUXI0lLtqxNxyagrgMWSDoReAq4APhAwzlPAmcB35d0LPCHwNa8SkeduCJiK/DG0V5vZtXVqYesI+KApBXAWqAHWBURmyRdXCtfCXwSuFnSo2Rdy8siYk9evZ45b2ZDZKtDdG6mVESsAdY0HFtZ93kH2VBTYU5cZjZE9shPtad4OnGZWYPqP/LjxGVmTYrMii+TE5eZDdHhu4rjYtIkrh1vy2/6Hn7Sb/IrWNW5WKxDpuRPYYnjX2pZdtbsx3OvvU9vGVVIE4W7imaWFK85b2bJCeCAW1xmlhp3Fc0sLeGuopkl5uBCglXmxGVmTdziMrOkjHAhwVJMmsT1r+/579zyT20e0TOeVgE9J702t/zxP289+W7hjz+Ye+1r1j06qpgmgkAcGPTgvJklxmNcZpaWcFfRzBLjMS4zS5ITl5klJRADHpw3s9R4cN7MkhIenK+Oqcp9o7cl6JAbXxz1tS/94lUdjGTiCScuM0uLH7I2swS5xWVmSYmAgUEnLjNLjO8qmllSAncVzSw5Hpw3swRFlB1BvgmTuAbfujC3/MzpP+hOINY1J8z49aivnX/vQAcjmXiq3lVs+0CSpFWSdkvaWHdspqR7JP289vfo8Q3TzLolu6s4pdBWliLffDOwpOHY5cB9EbEAuK+2b2YTRESxrSxtE1dE3A/sbTi8FLil9vkW4NzOhmVmZYpQoa0sox3jOjYidgJExE5Js1udKGk5sBxgOoeP8uvMrFuCcpNSEePeSY2IvojojYjeqRw63l9nZh0QBbeyjDZx7ZI0B6D2d3fnQjKzUgXEoAptRUhaIukJSVskDTseLmmxpA2SNkn6Xrs6R5u4VgMX1j5fCNw1ynrMrII6NcYlqQe4HjgbOAVYJumUhnOOAr4E/FVEvAF4f7t6245xSboNWAzMktQPXAlcDXxN0keBJ4t80Xj71XsOyy2f3ePxtdQccsLxueXvm7l61HUf9stncssn+yyvDt4xXARsiYitAJJuJ7u591jdOR8A7oyIJ7PvjrY9uLaJKyKWtSg6q921ZpaeET6rOEvS+rr9vojoq9ufC2yv2+8HTm+o42RgqqT/BY4ErouIW/O+dMLMnDezDgmgeOLaExG9OeXDVdTYnjsE+DOyxtBhwI8kPRARP2tVqROXmTXpYFexH5hftz8P2DHMOXsi4gXgBUn3A28EWiauar+DyMxKUOyOYsG7iuuABZJOlDQNuIDs5l69u4AzJR0i6XCyruTmvErd4jKzZh1qcUXEAUkrgLVAD7AqIjZJurhWvjIiNkv6DvAIMAjcGBEbW9fqxGVmjaKzq0NExBpgTcOxlQ37nwE+U7TOCZO4Dnn9c2O6/uXHj+pMINYx2z8/I7f8jEMHc8tvenZe68LfPDuakCYPr8dlZump9rOKTlxm1iy/MVs6Jy4zG2pk87hK4cRlZk285ryZpceJy8yS466imaVGbnGlYfb6it9GqaieWcfklu9678kty2ae35977fdOvqnNt0/PLb3h+nNbls3e9cM2dU9iISi4SGBZnLjMrJlbXGaWHCcuM0uOE5eZJcUTUM0sRb6raGbpceIys9S4xZWIl2bmr2KdvzLU2AyeeWpuefTkjzdsf0frN4Tve83+3GunTMt/EdfdZ/5HbvnUNkMh/zfQOrZ/3npe7rV7B/Pn1h0+JT/2Yx9svUZbxf9dls9jXGaWlKDymd2Jy8yaOXGZWWpU8SfgnLjMrJlbXGaWEoXvKppZinxX0cyS4xZXd7zy8tTc8sE2/0v85xXX5pavXrFwpCEVdtkxN+aWT2nzqqiXYl/Lsh0D+XOdvvj04tzyd9x7aW75UT+Zlls+5+5dLcv0q/z1uJ7efFhu+bE9+XPUYt2jueXWWtW7ivmzLgFJqyTtlrSx7thVkp6StKG2nTO+YZpZ10R2V7HIVpa2iQu4GVgyzPFrI2JhbVszTLmZpSoKbiVpm7gi4n5gbxdiMbOqSD1x5Vgh6ZFaV/LoVidJWi5pvaT1+3llDF9nZt1ycEpEu60so01cNwAnAQuBncA1rU6MiL6I6I2I3qm0fuDWzKyoUSWuiNgVEQMRMQh8GVjU2bDMrFQTsasoaU7d7nnAxlbnmlliErir2HYel6TbgMXALEn9wJXAYkkLyXLuNuDj4xdiMa//4E9yy9/w7ytyy+ef9lQnwxmR7+5u/e5BgKe/PS+3/JhNreczTfvOujbfnj8X6mTWt7k+X94ssqcue0vutacd+qPc8tufnzuKiKyQis/japu4ImLZMIfbvanTzBIlqj8BdcLMnDezDqp44hrLdAgzm4gKToUo2iqTtETSE5K2SLo857zTJA1Iel+7Op24zKzZYMGtDUk9wPXA2cApwDJJp7Q471PA2iLhOXGZWZMOtrgWAVsiYmtE7ANuB5YOc97fA3cAu4tU6sRlZs2Kz+OadfDJmNq2vKGmucD2uv3+2rHfkTSXbFrVyqLhTZrB+RP/Mf/WepXN4cmyQxgXh7/t6TFd/0/ffW9u+cn8eEz1T1ojm1y6JyJ6c8qHW5OpsfbPA5dFxIBUbAHDSZO4zKy4Dk6H6Afm1+3PA3Y0nNML3F5LWrOAcyQdiIhvtqrUicvMmnUuca0DFkg6EXgKuAD4wJCvijjx4GdJNwPfykta4MRlZsPo1OM8EXFA0gqyu4U9wKqI2CTp4lp54XGtek5cZjZUhx+gri00uqbh2LAJKyI+UqROJy4zG0IMP6JeJU5cZtas4o/8OHGZWRM/ZG02Tl57V8X/daWs4j+tE5eZDRXlLhJYhBOXmTVzi8vMUuMxLjNLjxOXmaXGLS4zS0tQaJHAMjlxmdkQflmGmaXJicvMUqOoduZy4jKzoTq8OsR4cOIysyYe4zKz5PiRHzNLj1tcZpaUEbyluixOXGbWLPXEJWk+cCtwHNl82r6IuE7STOC/gBOAbcD5EfHM+IVqk02P8t9X/MzJU3PLj/t2J6OZPFKYgFrkTdYHgE9ExB8DbwYukXQKcDlwX0QsAO6r7ZvZBKDBKLSVpW3iioidEfFw7fNzwGayV2gvBW6pnXYLcO44xWhm3RQj2EoyojEuSScApwIPAsdGxE7Ikpuk2Z0Pz8zKMGGmQ0g6ArgDuDQinq29LrvIdcuB5QDTOXw0MZpZt02AMS4kTSVLWl+JiDtrh3dJmlMrnwPsHu7aiOiLiN6I6J3KoZ2I2czGmaLYVpa2iUtZ0+omYHNEfK6uaDVwYe3zhcBdnQ/PzLougIhiW0mKdBXPAD4EPCppQ+3YFcDVwNckfRR4Enj/uERok9ZAtBloKdRfsNFIfowrIn5A6zdyn9XZcMysbCnM4/LMeTMbquRuYBFOXGbWxC0uM0uPE5eZpcYtLjNLSwAD1c5cTlxm1sQtLrNx8uJpL5YdwsTVwbuKkpYA1wE9wI0RcXVD+d8Al9V2nwf+LiJ+mlenE5eZNelUi0tSD3A98E6gH1gnaXVEPFZ32i+BP4+IZySdDfQBp+fV67nHZjZUZ5e1WQRsiYitEbEPuJ1sSazff13ED+sWIX0AmNeuUre4zGwIASo+OD9L0vq6/b6I6Kvbnwtsr9vvJ7819VGg7dq1Tlxm1mQEb7LeExG9eVUNc2zYyiW9nSxxvbXdlzpxmdlQnV3dtB+YX7c/D9jReJKkPwVuBM6OiF+3q9RjXGbWoOCSNsVaZeuABZJOlDQNuIBsSazfkXQ8cCfwoYj4WZFK3eIysyaduqsYEQckrQDWkk2HWBURmyRdXCtfCfwLcAzwpdrKygfadD+duKy62r2ezMZRB+dxRcQaYE3DsZV1nz8GfGwkdTpxmdlQMaK7iqVw4jKzZtXOW05cZtZsBNMhSuHEZWbNnLjMLCkBpP6yDDObXES4q2hmCRqsdpPLictK88q9r84tH1hY7X88E5a7imaWIncVzSw9Tlxmlha/ENbMUuO3/JhZijzGZWbpceIys6QEMJh44pI0H7gVOI5sdkdfRFwn6Srgb4Gna6deUVt3x6yQ4679YW75Ode+Kbf8dWzoYDT2exNjcP4A8ImIeFjSkcBDku6plV0bEZ8dv/DMrBSpJ66I2AnsrH1+TtJmslcOmdlEFMBAtafOj2htXEknAKcCD9YOrZD0iKRVko5ucc1ySeslrd/PK2OL1sy6ICAGi20lKZy4JB0B3AFcGhHPAjcAJwELyVpk1wx3XUT0RURvRPRO5dCxR2xm469zb/kZF4XuKkqaSpa0vhIRdwJExK668i8D3xqXCM2suxK4q9i2xaXsfUE3AZsj4nN1x+fUnXYesLHz4ZlZKSZAi+sM4EPAo5I21I5dASyTtJAsP28DPj4O8ZlZGSbAXcUfABqmyHO2zCaiCBgYKDuKXJ45b2bNUm9xmdkk5MRlZmmJyt9VdOIys6ECosTJpUU4cZlZs4o/8uPEZWZDRfj1ZGaWIA/Om1lqwi0uM0vLxFhI0MwmkwQesnbiMrMhAoiKP/IzooUEzWwSiM4uJChpiaQnJG2RdPkw5ZL0hVr5I5LyXzaAW1xmNozoUFdRUg9wPfBOoB9YJ2l1RDxWd9rZwILadjrZIqWn59XrFpeZNetci2sRsCUitkbEPuB2YGnDOUuBWyPzAHBUw3p/Tbra4nqOZ/bcG1//Vd2hWcCebsYwAlWNrapxgWMbrU7G9tqxVvAcz6y9N74+q+Dp0yWtr9vvi4i+uv25wPa6/X6aW1PDnTOX2kt6htPVxBURr67fl7Q+Inq7GUNRVY2tqnGBYxutqsUWEUs6WN1wa/k19kOLnDOEu4pmNp76gfl1+/OAHaM4ZwgnLjMbT+uABZJOlDQNuABY3XDOauDDtbuLbwZ+W3ufa0tl31Xsa39KaaoaW1XjAsc2WlWObUwi4oCkFcBaoAdYFRGbJF1cK19Jtgz8OcAW4EXgonb1Kio+td/MrJG7imaWHCcuM0tOKYmr3SMAZZK0TdKjkjY0zE8pI5ZVknZL2lh3bKakeyT9vPb36ArFdpWkp2q/3QZJ55QU23xJ35W0WdImSf9QO17qb5cTVyV+t5R0fYyr9gjAz6h7BABY1vAIQGkkbQN6I6L0yYqS3gY8Tzar+E9qxz4N7I2Iq2tJ/+iIuKwisV0FPB8Rn+12PA2xzQHmRMTDko4EHgLOBT5Cib9dTlznU4HfLSVltLiKPAJgQETcD+xtOLwUuKX2+Ray/+N3XYvYKiEidkbEw7XPzwGbyWZil/rb5cRlI1RG4mo1vb8qArhb0kOSlpcdzDCOPTjHpfZ3dsnxNFpRe8J/VVnd2HqSTgBOBR6kQr9dQ1xQsd+t6spIXCOe3t9lZ0TEm8ieWL+k1iWyYm4ATgIWkj1ndk2ZwUg6ArgDuDQini0zlnrDxFWp3y0FZSSuEU/v76aI2FH7uxv4BlnXtkp2HXxyvvZ3d8nx/E5E7IqIgcheyvdlSvztJE0lSw5fiYg7a4dL/+2Gi6tKv1sqykhcRR4BKIWkGbVBUyTNAN4FbMy/qutWAxfWPl8I3FViLEM0LEVyHiX9dpIE3ARsjojP1RWV+tu1iqsqv1tKSpk5X7vd+3l+/wjAv3U9iGFIeh1ZKwuyx6G+WmZskm4DFpMte7ILuBL4JvA14HjgSeD9EdH1QfIWsS0m6+4EsA34eLtnzsYptrcC3wceBQ4uGnUF2XhSab9dTlzLqMDvlhI/8mNmyfHMeTNLjhOXmSXHicvMkuPEZWbJceIys+Q4cZlZcpy4zCw5/w+z+CzRYjVtjgAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ground truth: 4\n"
     ]
    }
   ],
   "source": [
    "# Sample visualization\n",
    "n = 2  # slice\n",
    "plt.imshow(X[n][0], interpolation='nearest')\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "print(\"Ground truth: {}\".format(y[n]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df2ff06d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "c4eeab30",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1\n",
      "-------------------------------\n",
      "loss: 2.303921  [    0/60000]\n",
      "loss: 0.242171  [ 6400/60000]\n",
      "loss: 0.143744  [12800/60000]\n",
      "loss: 0.140093  [19200/60000]\n",
      "loss: 0.079750  [25600/60000]\n",
      "loss: 0.170518  [32000/60000]\n",
      "loss: 0.202803  [38400/60000]\n",
      "loss: 0.204086  [44800/60000]\n",
      "loss: 0.306435  [51200/60000]\n",
      "loss: 0.076393  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 96.7%, Avg loss: 0.103248 \n",
      "\n",
      "Epoch 2\n",
      "-------------------------------\n",
      "loss: 0.090938  [    0/60000]\n",
      "loss: 0.094877  [ 6400/60000]\n",
      "loss: 0.042177  [12800/60000]\n",
      "loss: 0.102859  [19200/60000]\n",
      "loss: 0.028173  [25600/60000]\n",
      "loss: 0.041681  [32000/60000]\n",
      "loss: 0.079733  [38400/60000]\n",
      "loss: 0.081660  [44800/60000]\n",
      "loss: 0.191324  [51200/60000]\n",
      "loss: 0.055518  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 97.5%, Avg loss: 0.094891 \n",
      "\n",
      "Epoch 3\n",
      "-------------------------------\n",
      "loss: 0.098737  [    0/60000]\n",
      "loss: 0.013506  [ 6400/60000]\n",
      "loss: 0.025344  [12800/60000]\n",
      "loss: 0.124207  [19200/60000]\n",
      "loss: 0.055090  [25600/60000]\n",
      "loss: 0.050356  [32000/60000]\n",
      "loss: 0.004770  [38400/60000]\n",
      "loss: 0.006831  [44800/60000]\n",
      "loss: 0.034768  [51200/60000]\n",
      "loss: 0.036470  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 97.6%, Avg loss: 0.096849 \n",
      "\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "epochs = 3\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train(train_dataloader, model, loss_fn, optimizer)\n",
    "    test(test_dataloader, model, loss_fn)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6842317",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('base')",
   "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.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "7d658edeaff6acdc8b4336631f180adfdc2b0716532322f3654722bfb2182602"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}