ResNet56_From_Scratch_PyTorch.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## ResNet Paper Implementation from Scratch with PyTorch on CIFAR-10 Dataset\n",
        "\n",
        "---\n",
        "\n",
        "The below comments are taken from the original paper from implementing ResNet on CIFAR-10 Dataset. And I will follow this structure for this implementation of ResNet on CIFAR10.\n",
        "\n",
        "\"We conducted more studies on the CIFAR-10 dataset\n",
        "which consists of 50k training images and 10k test-\n",
        "ing images in 10 classes.\n",
        "\n",
        " The network inputs are 32×32 images, with\n",
        "the per-pixel mean subtracted. The first layer is 3×3 convolutions. Then we use a stack of 6n layers with 3×3 convolutions on the feature maps of sizes {32, 16, 8} respectively,\n",
        "with 2n layers for each feature map size. The numbers of\n",
        "filters are {16, 32, 64} respectively.\n",
        "\n",
        "The subsampling is performed by convolutions with a stride of 2. The network ends\n",
        "with a global average pooling, a 10-way fully-connected\n",
        "layer, and softmax.\n",
        "\n",
        "There are totally 6n+2 stacked weighted layers.\n",
        "\n",
        "We compare n = {3, 5, 7, 9}, leading to 20, 32, 44, and 56-layer networks.\n",
        "\n",
        "When shortcut connections are used, they are connected\n",
        "to the pairs of 3×3 layers (totally 3n shortcuts).\n",
        "\n",
        "On this dataset we use identity shortcuts in all cases (i.e., option A), so our residual models have exactly the same depth, width, and number of parameters as the plain counterparts.\"\n",
        "---\n",
        "\n",
        "For this example here in this file, I have used n=9. So my ResNet blocks are like [9, 9, 9]\n",
        "\n",
        "And thats why the total layers are 56 (i.e. 9 * 6 + 2)\n",
        "\n",
        "![Imgur](https://imgur.com/ifD8qbd.png)\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "There are two kinds of shortcut module in the implementation of ResNet.\n",
        "\n",
        "![Imgur](https://imgur.com/DZPWHHo.png)\n",
        "\n",
        "\n",
        "### The identity_block\n",
        "\n",
        "The first one is an identity block which has no convolution layer at shortcut.  In this case, input has the same dimensions as output.\n",
        "\n",
        "**The identity_block** is the one with the straight line. It consists of three convolution layers (with Batch Norm and a ReLU). The input of the block is added to the last one right before the final activation function. The connection from the input to the add operation is called a skip connection (or a shortcut as keras calls it).\n",
        "\n",
        "\n",
        "### The conv_block\n",
        "\n",
        "The other is convolution block, which has convolution layer at shortcut. In this case, the input dimensions are smaller than the output dimensions.\n",
        "\n",
        "\n",
        "**The conv_block** is the one with the dotted line. It two consists the three convolution layers (+BN+ReLU); these are different than the previous, though. The difference is that this time the skip connection passes through an independent convolution layer before being added to the output of the third convolution layer. This layer has a kernel of 1×1 and uses strides of 2, which has the effect that changing the dimension of its input.\n",
        "\n",
        "The convolutional block is used whenever you need to change the dimension in order for the output and input dimensions to match. That being said, how do you change the dimension of a certain volume using convolutions? Well, you change the stride.\n",
        "For any given convolution operation, assuming a square input, the dimension of the output volume can be obtained through the formula (n+2p-f)/s +1, where n is the input dimension, p is your zero-padding, f the filter dimension and s is the stride. By increasing the stride you're effectively reducing the dimension of your shortcut's output volume, and thus, it can be used in such a way as to make sure that the dimensions of your shortcut and lower paths will match in order for the final sum to be performed.\n",
        "\n",
        "---\n",
        "\n",
        "### Making sense of Residual intutively\n",
        "\n",
        "![Imgur](https://imgur.com/Ah27NSU.png)\n",
        "\n",
        "\n",
        "In traditional neural networks, each layer feeds into the next layer. In a network with residual blocks, each layer feeds into the next layer and directly into the layers about 2–3 hops away. That’s it.\n",
        "\n",
        "We know neural networks are universal function approximators and that the accuracy increases with increasing the number of layers. So, if neural networks were universal function approximators, then they should have been able to learn any simplex or complex function. But it turns out that, thanks to some problems like vanishing gradients and the curse of dimensionality, if we have sufficiently deep networks, it may not be able to learn simple functions like an identity function. Now, this is clearly undesirable.\n",
        "\n",
        "Also, if we keep increasing the number of layers, we will see that the accuracy will saturate at one point and eventually degrade. And, this is usually not caused due to overfitting. So, it might seem that the shallower networks are learning better than their deeper counterparts, and this is quite counter-intuitive. But this is what is seen in practice and is popularly known as the degradation problem.\n",
        "\n",
        "\n",
        "#### In the degradation problem, we know that shallower networks perform better than the deeper counterparts. So, why not skip these extra layers and at least match the accuracy of the shallow sub-networks. But, how do you skip layers?\n",
        "\n",
        "#### You can skip the training of few layers using skip connections or residual connections. This is what we see in the image above. In fact, if you look closely, we can directly learn an identity function by relying on skip connections only. This is the exact reason why skip connections are also called identity shortcut connections too.\n",
        "\n",
        "---\n",
        "\n",
        "### Lets start the implementation\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "9g7Qx3Sb7-tD"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import shutil\n",
        "from collections import OrderedDict\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torchvision import transforms, datasets\n",
        "from torchsummary import summary\n",
        "from torch.utils.data import Dataset, DataLoader, random_split"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9mPnawhC7-tF",
        "outputId": "646b4228-0910-48d0-c56a-e0cf599b801b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5ZL8T3Do8J5s",
        "outputId": "81edc139-34c1-4dbe-cb6e-06d1e48ef1a6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Sat Mar  5 13:11:50 2022       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla P100-PCIE...  Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   35C    P0    26W / 250W |      0MiB / 16280MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "pwngtd6i7-tL"
      },
      "outputs": [],
      "source": [
        "class LambdaLayer(nn.Module):\n",
        "    \n",
        "    def __init__(self, lambd):\n",
        "        super(LambdaLayer, self).__init__()\n",
        "        self.lambd = lambd\n",
        "    \n",
        "    def forward(self, x):\n",
        "        return self.lambd(x)\n",
        "\n",
        "class BasicConvBlock(nn.Module):\n",
        "    \n",
        "    ''' The BasicConvBlock takes an input with in_channels, applies some blocks of convolutional layers \n",
        "    to reduce it to out_channels and sum it up to the original input. \n",
        "    If their sizes mismatch, then the input goes into an identity. \n",
        "    \n",
        "    Basically The BasicConvBlock will implement the regular basic Conv Block + \n",
        "    the shortcut block that does the dimension matching job (option A or B) when dimension changes between 2 blocks\n",
        "    '''\n",
        "    \n",
        "    def __init__(self, in_channels, out_channels, stride=1, option='A'):\n",
        "        super(BasicConvBlock, self).__init__()\n",
        "        \n",
        "        self.features = nn.Sequential(OrderedDict([\n",
        "            ('conv1', nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)),\n",
        "            ('bn1', nn.BatchNorm2d(out_channels)),\n",
        "            ('act1', nn.ReLU()),\n",
        "            ('conv2', nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)),\n",
        "            ('bn2', nn.BatchNorm2d(out_channels))\n",
        "        ]))\n",
        "\n",
        "        self.shortcut = nn.Sequential()\n",
        "        \n",
        "        '''  When input and output spatial dimensions don't match, we have 2 options, with stride:\n",
        "            - A) Use identity shortcuts with zero padding to increase channel dimension.    \n",
        "            - B) Use 1x1 convolution to increase channel dimension (projection shortcut).\n",
        "         '''\n",
        "        if stride != 1 or in_channels != out_channels:\n",
        "            if option == 'A':\n",
        "                # Use identity shortcuts with zero padding to increase channel dimension.\n",
        "                pad_to_add = out_channels//4\n",
        "                ''' ::2 is doing the job of stride = 2\n",
        "                F.pad apply padding to (W,H,C,N).\n",
        "                \n",
        "                The padding lengths are specified in reverse order of the dimensions,\n",
        "                F.pad(x[:, :, ::2, ::2], (0,0, 0,0, pad,pad, 0,0))\n",
        "\n",
        "                [width_beginning, width_end, height_beginning, height_end, channel_beginning, channel_end, batchLength_beginning, batchLength_end ]\n",
        "\n",
        "                '''\n",
        "                self.shortcut = LambdaLayer(lambda x:\n",
        "                            F.pad(x[:, :, ::2, ::2], (0,0, 0,0, pad_to_add, pad_to_add, 0,0)))\n",
        "            if option == 'B':\n",
        "                self.shortcut = nn.Sequential(OrderedDict([\n",
        "                    ('s_conv1', nn.Conv2d(in_channels, 2*out_channels, kernel_size=1, stride=stride, padding=0, bias=False)),\n",
        "                    ('s_bn1', nn.BatchNorm2d(2*out_channels))\n",
        "                ]))\n",
        "        \n",
        "    def forward(self, x):\n",
        "        out = self.features(x)\n",
        "        # sum it up with shortcut layer\n",
        "        out += self.shortcut(x)\n",
        "        out = F.relu(out)\n",
        "        return out"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n",
        "\n",
        "### Explanations on using Option A and B in below code\n",
        "\n",
        "```py\n",
        "\n",
        "if stride != 1 or in_channels != out_channels:\n",
        "            if option == 'A':\n",
        "                pad = out_channels//4\n",
        "                # ::2 replace the stride 2 + F.pad apply padding to (W,H,C,N).\n",
        "                self.shortcut = LambdaLayer(lambda x:\n",
        "                            F.pad(x[:, :, ::2, ::2], (0,0, 0,0, pad,pad, 0,0)))\n",
        "            if option == 'B':\n",
        "                self.shortcut = nn.Sequential(OrderedDict([\n",
        "                    ('s_conv1', nn.Conv2d(in_channels, 2*out_channels, kernel_size=1, stride=stride, padding=0, bias=False)),\n",
        "                    ('s_bn1', nn.BatchNorm2d(2*out_channels))\n",
        "                ]))\n",
        "\n",
        "```\n",
        "\n",
        "The above situation of using Option A or B arises, because, if the number of filters in the input layer does not match the number of filters in the last convolutional layer of the module, then we will get an error.\n",
        "\n",
        "One solution is to use a 1×1 convolution layer, often referred to as a projection layer, to either increase the number of filters for the input layer or reduce the number of filters for the last convolutional layer in the module. The former solution makes more sense, and is the approach proposed in the paper, referred to as a projection shortcut.\n",
        "\n",
        "When the dimensions increase, we consider two options:\n",
        "\n",
        "- (A) The shortcut still performs identity mapping, with extra zero entries padded for increasing dimensions. This option introduces no extra parameter;\n",
        "\n",
        "- (B) The projection shortcut is used to match dimensions (done by 1×1 convolutions).\n",
        "\n",
        "\n",
        "This is pretty much what the the original Paper talks about as well.\n",
        "\n",
        "#### We use identity shortcuts when input and output channel dimensions are the same.\n",
        "\n",
        "#### Otherwise, When input and output spatial dimensions don't match, we have 2 options, with stride:\n",
        "\n",
        "    - A) Use identity shortcuts with zero padding to increase channel dimension.\n",
        "\n",
        "    - B) Use 1x1 convolution to increase channel dimension (projection shortcut).\n",
        "\n",
        "-----------------------\n",
        "\n",
        "### Understanding `F.pad` on a 4-D Tensor and the following line\n",
        "\n",
        "### `F.pad(x[:, :, ::2, ::2], (0,0, 0,0, pad,pad, 0,0)))`\n",
        "\n",
        "https://stackoverflow.com/a/61945903/1902852\n",
        "\n",
        "The padding lengths are specified in reverse order of the dimensions, where every dimension has two values, one for the padding at the beginning and one for the padding at the end.\n",
        "\n",
        "For an image with the dimensions `[channels, height, width]` the padding is given as:\n",
        "\n",
        "`[width_beginning, width_end, height_beginning, height_end, channels_beginning, channels_end]`,\n",
        "\n",
        "which can be reworded to\n",
        "\n",
        "`[left, right, top, bottom]`\n",
        "\n",
        "Therefore the code above pads the images to the right and bottom. The channels are left out, because they are not being padded, which also means that the same padding could be directly applied to the masks.\n",
        "\n",
        "So the below line means\n",
        "\n",
        "`F.pad(x[:, :, ::2, ::2], (0,0, 0,0, pad,pad, 0,0))`\n",
        "\n",
        "\n",
        "`[width_beginning, width_end, height_beginning, height_end, channel_beginning, channel_end, batchLength_beginning, batchLength_end ]`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "6jvgyHXI7-tM"
      },
      "outputs": [],
      "source": [
        "\n",
        "class ResNet(nn.Module):\n",
        "    \"\"\"\n",
        "        ResNet-56 architecture for CIFAR-10 Dataset of shape 32*32*3\n",
        "    \"\"\"\n",
        "    def __init__(self, block_type, num_blocks):\n",
        "        super(ResNet, self).__init__()\n",
        "        \n",
        "        self.in_channels = 16\n",
        "        \n",
        "        self.conv0 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "        self.bn0 = nn.BatchNorm2d(16)\n",
        "        \n",
        "        self.block1 = self.__build_layer(block_type, 16, num_blocks[0], starting_stride=1)\n",
        "        \n",
        "        self.block2 = self.__build_layer(block_type, 32, num_blocks[1], starting_stride=2)\n",
        "        \n",
        "        self.block3 = self.__build_layer(block_type, 64, num_blocks[2], starting_stride=2)\n",
        "        \n",
        "        self.avgpool = nn.AdaptiveAvgPool2d((1,1))\n",
        "        self.linear = nn.Linear(64, 10)\n",
        "    \n",
        "    def __build_layer(self, block_type, out_channels, num_blocks, starting_stride):\n",
        "        \n",
        "        strides_list_for_current_block = [starting_stride] + [1]*(num_blocks-1)\n",
        "        ''' Above line will generate an array whose first element is starting_stride\n",
        "        And it will have (num_blocks-1) more elements each of value 1\n",
        "         '''\n",
        "        # print('strides_list_for_current_block ', strides_list_for_current_block)\n",
        "        \n",
        "        layers = []\n",
        "        \n",
        "        for stride in strides_list_for_current_block:\n",
        "            layers.append(block_type(self.in_channels, out_channels, stride))\n",
        "            self.in_channels = out_channels\n",
        "        \n",
        "        return nn.Sequential(*layers)\n",
        "    \n",
        "    def forward(self, x):\n",
        "        out = F.relu(self.bn0(self.conv0(x)))\n",
        "        out = self.block1(out)\n",
        "        out = self.block2(out)        \n",
        "        out = self.block3(out)\n",
        "        out = self.avgpool(out)\n",
        "        out = torch.flatten(out, 1)\n",
        "        out = self.linear(out)\n",
        "        return out"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### _build_layer() method\n",
        "\n",
        "In ResNet Every layer downsamples the input at the start using stride equals to 2 i.e for 1st convolutional layer in 1st block of a layer.\n",
        "\n",
        "If we look at the first operation of each layer, we see that the stride used at that first one is 2, instead of 1 like for the rest of them.\n",
        "\n",
        "This is because, here in ResNet, reduction between layers is achieved by an increase on the stride, from 1 to 2, at the first convolution of each layer; instead of by a pooling operation, which we are used to see as down samplers.\n",
        "\n",
        "Quoting from Paper\n",
        "\n",
        "\" For both options, when the shortcuts go across feature maps of two sizes, they are performed with a stride of 2.\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "pnvR0RuC7-tN"
      },
      "outputs": [],
      "source": [
        "def ResNet56():\n",
        "    return ResNet(block_type=BasicConvBlock, num_blocks=[9,9,9])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "47k9cNnn7-tN",
        "outputId": "9832007a-4e1f-4c88-c4c0-931f0af67bf6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "----------------------------------------------------------------\n",
            "        Layer (type)               Output Shape         Param #\n",
            "================================================================\n",
            "            Conv2d-1           [-1, 16, 32, 32]             432\n",
            "       BatchNorm2d-2           [-1, 16, 32, 32]              32\n",
            "            Conv2d-3           [-1, 16, 32, 32]           2,304\n",
            "       BatchNorm2d-4           [-1, 16, 32, 32]              32\n",
            "              ReLU-5           [-1, 16, 32, 32]               0\n",
            "            Conv2d-6           [-1, 16, 32, 32]           2,304\n",
            "       BatchNorm2d-7           [-1, 16, 32, 32]              32\n",
            "    BasicConvBlock-8           [-1, 16, 32, 32]               0\n",
            "            Conv2d-9           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-10           [-1, 16, 32, 32]              32\n",
            "             ReLU-11           [-1, 16, 32, 32]               0\n",
            "           Conv2d-12           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-13           [-1, 16, 32, 32]              32\n",
            "   BasicConvBlock-14           [-1, 16, 32, 32]               0\n",
            "           Conv2d-15           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-16           [-1, 16, 32, 32]              32\n",
            "             ReLU-17           [-1, 16, 32, 32]               0\n",
            "           Conv2d-18           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-19           [-1, 16, 32, 32]              32\n",
            "   BasicConvBlock-20           [-1, 16, 32, 32]               0\n",
            "           Conv2d-21           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-22           [-1, 16, 32, 32]              32\n",
            "             ReLU-23           [-1, 16, 32, 32]               0\n",
            "           Conv2d-24           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-25           [-1, 16, 32, 32]              32\n",
            "   BasicConvBlock-26           [-1, 16, 32, 32]               0\n",
            "           Conv2d-27           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-28           [-1, 16, 32, 32]              32\n",
            "             ReLU-29           [-1, 16, 32, 32]               0\n",
            "           Conv2d-30           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-31           [-1, 16, 32, 32]              32\n",
            "   BasicConvBlock-32           [-1, 16, 32, 32]               0\n",
            "           Conv2d-33           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-34           [-1, 16, 32, 32]              32\n",
            "             ReLU-35           [-1, 16, 32, 32]               0\n",
            "           Conv2d-36           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-37           [-1, 16, 32, 32]              32\n",
            "   BasicConvBlock-38           [-1, 16, 32, 32]               0\n",
            "           Conv2d-39           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-40           [-1, 16, 32, 32]              32\n",
            "             ReLU-41           [-1, 16, 32, 32]               0\n",
            "           Conv2d-42           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-43           [-1, 16, 32, 32]              32\n",
            "   BasicConvBlock-44           [-1, 16, 32, 32]               0\n",
            "           Conv2d-45           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-46           [-1, 16, 32, 32]              32\n",
            "             ReLU-47           [-1, 16, 32, 32]               0\n",
            "           Conv2d-48           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-49           [-1, 16, 32, 32]              32\n",
            "   BasicConvBlock-50           [-1, 16, 32, 32]               0\n",
            "           Conv2d-51           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-52           [-1, 16, 32, 32]              32\n",
            "             ReLU-53           [-1, 16, 32, 32]               0\n",
            "           Conv2d-54           [-1, 16, 32, 32]           2,304\n",
            "      BatchNorm2d-55           [-1, 16, 32, 32]              32\n",
            "   BasicConvBlock-56           [-1, 16, 32, 32]               0\n",
            "           Conv2d-57           [-1, 32, 16, 16]           4,608\n",
            "      BatchNorm2d-58           [-1, 32, 16, 16]              64\n",
            "             ReLU-59           [-1, 32, 16, 16]               0\n",
            "           Conv2d-60           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-61           [-1, 32, 16, 16]              64\n",
            "      LambdaLayer-62           [-1, 32, 16, 16]               0\n",
            "   BasicConvBlock-63           [-1, 32, 16, 16]               0\n",
            "           Conv2d-64           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-65           [-1, 32, 16, 16]              64\n",
            "             ReLU-66           [-1, 32, 16, 16]               0\n",
            "           Conv2d-67           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-68           [-1, 32, 16, 16]              64\n",
            "   BasicConvBlock-69           [-1, 32, 16, 16]               0\n",
            "           Conv2d-70           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-71           [-1, 32, 16, 16]              64\n",
            "             ReLU-72           [-1, 32, 16, 16]               0\n",
            "           Conv2d-73           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-74           [-1, 32, 16, 16]              64\n",
            "   BasicConvBlock-75           [-1, 32, 16, 16]               0\n",
            "           Conv2d-76           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-77           [-1, 32, 16, 16]              64\n",
            "             ReLU-78           [-1, 32, 16, 16]               0\n",
            "           Conv2d-79           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-80           [-1, 32, 16, 16]              64\n",
            "   BasicConvBlock-81           [-1, 32, 16, 16]               0\n",
            "           Conv2d-82           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-83           [-1, 32, 16, 16]              64\n",
            "             ReLU-84           [-1, 32, 16, 16]               0\n",
            "           Conv2d-85           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-86           [-1, 32, 16, 16]              64\n",
            "   BasicConvBlock-87           [-1, 32, 16, 16]               0\n",
            "           Conv2d-88           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-89           [-1, 32, 16, 16]              64\n",
            "             ReLU-90           [-1, 32, 16, 16]               0\n",
            "           Conv2d-91           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-92           [-1, 32, 16, 16]              64\n",
            "   BasicConvBlock-93           [-1, 32, 16, 16]               0\n",
            "           Conv2d-94           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-95           [-1, 32, 16, 16]              64\n",
            "             ReLU-96           [-1, 32, 16, 16]               0\n",
            "           Conv2d-97           [-1, 32, 16, 16]           9,216\n",
            "      BatchNorm2d-98           [-1, 32, 16, 16]              64\n",
            "   BasicConvBlock-99           [-1, 32, 16, 16]               0\n",
            "          Conv2d-100           [-1, 32, 16, 16]           9,216\n",
            "     BatchNorm2d-101           [-1, 32, 16, 16]              64\n",
            "            ReLU-102           [-1, 32, 16, 16]               0\n",
            "          Conv2d-103           [-1, 32, 16, 16]           9,216\n",
            "     BatchNorm2d-104           [-1, 32, 16, 16]              64\n",
            "  BasicConvBlock-105           [-1, 32, 16, 16]               0\n",
            "          Conv2d-106           [-1, 32, 16, 16]           9,216\n",
            "     BatchNorm2d-107           [-1, 32, 16, 16]              64\n",
            "            ReLU-108           [-1, 32, 16, 16]               0\n",
            "          Conv2d-109           [-1, 32, 16, 16]           9,216\n",
            "     BatchNorm2d-110           [-1, 32, 16, 16]              64\n",
            "  BasicConvBlock-111           [-1, 32, 16, 16]               0\n",
            "          Conv2d-112             [-1, 64, 8, 8]          18,432\n",
            "     BatchNorm2d-113             [-1, 64, 8, 8]             128\n",
            "            ReLU-114             [-1, 64, 8, 8]               0\n",
            "          Conv2d-115             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-116             [-1, 64, 8, 8]             128\n",
            "     LambdaLayer-117             [-1, 64, 8, 8]               0\n",
            "  BasicConvBlock-118             [-1, 64, 8, 8]               0\n",
            "          Conv2d-119             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-120             [-1, 64, 8, 8]             128\n",
            "            ReLU-121             [-1, 64, 8, 8]               0\n",
            "          Conv2d-122             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-123             [-1, 64, 8, 8]             128\n",
            "  BasicConvBlock-124             [-1, 64, 8, 8]               0\n",
            "          Conv2d-125             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-126             [-1, 64, 8, 8]             128\n",
            "            ReLU-127             [-1, 64, 8, 8]               0\n",
            "          Conv2d-128             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-129             [-1, 64, 8, 8]             128\n",
            "  BasicConvBlock-130             [-1, 64, 8, 8]               0\n",
            "          Conv2d-131             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-132             [-1, 64, 8, 8]             128\n",
            "            ReLU-133             [-1, 64, 8, 8]               0\n",
            "          Conv2d-134             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-135             [-1, 64, 8, 8]             128\n",
            "  BasicConvBlock-136             [-1, 64, 8, 8]               0\n",
            "          Conv2d-137             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-138             [-1, 64, 8, 8]             128\n",
            "            ReLU-139             [-1, 64, 8, 8]               0\n",
            "          Conv2d-140             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-141             [-1, 64, 8, 8]             128\n",
            "  BasicConvBlock-142             [-1, 64, 8, 8]               0\n",
            "          Conv2d-143             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-144             [-1, 64, 8, 8]             128\n",
            "            ReLU-145             [-1, 64, 8, 8]               0\n",
            "          Conv2d-146             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-147             [-1, 64, 8, 8]             128\n",
            "  BasicConvBlock-148             [-1, 64, 8, 8]               0\n",
            "          Conv2d-149             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-150             [-1, 64, 8, 8]             128\n",
            "            ReLU-151             [-1, 64, 8, 8]               0\n",
            "          Conv2d-152             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-153             [-1, 64, 8, 8]             128\n",
            "  BasicConvBlock-154             [-1, 64, 8, 8]               0\n",
            "          Conv2d-155             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-156             [-1, 64, 8, 8]             128\n",
            "            ReLU-157             [-1, 64, 8, 8]               0\n",
            "          Conv2d-158             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-159             [-1, 64, 8, 8]             128\n",
            "  BasicConvBlock-160             [-1, 64, 8, 8]               0\n",
            "          Conv2d-161             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-162             [-1, 64, 8, 8]             128\n",
            "            ReLU-163             [-1, 64, 8, 8]               0\n",
            "          Conv2d-164             [-1, 64, 8, 8]          36,864\n",
            "     BatchNorm2d-165             [-1, 64, 8, 8]             128\n",
            "  BasicConvBlock-166             [-1, 64, 8, 8]               0\n",
            "AdaptiveAvgPool2d-167             [-1, 64, 1, 1]               0\n",
            "          Linear-168                   [-1, 10]             650\n",
            "================================================================\n",
            "Total params: 853,018\n",
            "Trainable params: 853,018\n",
            "Non-trainable params: 0\n",
            "----------------------------------------------------------------\n",
            "Input size (MB): 0.01\n",
            "Forward/backward pass size (MB): 12.16\n",
            "Params size (MB): 3.25\n",
            "Estimated Total Size (MB): 15.42\n",
            "----------------------------------------------------------------\n"
          ]
        }
      ],
      "source": [
        "model = ResNet56()\n",
        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "# device = 'cpu'\n",
        "model.to(device)\n",
        "summary(model, (3, 32, 32))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QjPLm6BENFbq"
      },
      "source": [
        "## Loading CIFAR-10 Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "Us3eora2NFbs"
      },
      "outputs": [],
      "source": [
        "def dataloader_cifar():\n",
        "    transform = transforms.Compose([transforms.ToTensor(),\n",
        "                                    transforms.Normalize(mean=[0.5], std=[0.5])])\n",
        "            \n",
        "    # Input Data in Local Machine\n",
        "    # train_dataset = datasets.CIFAR10('../input_data', train=True, download=True, transform=transform)\n",
        "    # test_dataset = datasets.CIFAR10('../input_data', train=False, download=True, transform=transform)\n",
        "    \n",
        "    # Input Data in Google Drive\n",
        "    train_dataset = datasets.CIFAR10('/content/drive/MyDrive/All_Datasets/CIFAR10', train=True, download=True, transform=transform)\n",
        "    test_dataset = datasets.CIFAR10('/content/drive/MyDrive/All_Datasets/CIFAR10', train=False, download=True, transform=transform)\n",
        "\n",
        "    # Split dataset into training set and validation set.\n",
        "    train_dataset, val_dataset = random_split(train_dataset, (45000, 5000))\n",
        "    \n",
        "    print(\"Image shape of a random sample image : {}\".format(train_dataset[0][0].numpy().shape), end = '\\n\\n')\n",
        "    \n",
        "    print(\"Training Set:   {} images\".format(len(train_dataset)))\n",
        "    print(\"Validation Set:   {} images\".format(len(val_dataset)))\n",
        "    print(\"Test Set:       {} images\".format(len(test_dataset)))\n",
        "    \n",
        "    BATCH_SIZE = 32\n",
        "\n",
        "    # Generate dataloader\n",
        "    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
        "    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
        "    test_loader = DataLoader(test_dataset, batch_size=10000, shuffle=True)\n",
        "    \n",
        "    return train_loader, val_loader, test_loader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4E--f56MNFbs",
        "outputId": "771fcf9c-e61d-4613-cab6-0e1a97844207"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Files already downloaded and verified\n",
            "Files already downloaded and verified\n",
            "Image shape of a random sample image : (3, 32, 32)\n",
            "\n",
            "Training Set:   45000 images\n",
            "Validation Set:   5000 images\n",
            "Test Set:       10000 images\n"
          ]
        }
      ],
      "source": [
        "train_loader, val_loader, test_loader = dataloader_cifar()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tOv_8vpcNFbt"
      },
      "source": [
        "## Start Actual Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "ICxcakbK8DRP"
      },
      "outputs": [],
      "source": [
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.Adam(model.parameters(), lr=0.01)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "sB_vMJgl8Ei5"
      },
      "outputs": [],
      "source": [
        "def train_model():\n",
        "    EPOCHS = 15\n",
        "    train_samples_num = 45000\n",
        "    val_samples_num = 5000\n",
        "    train_costs, val_costs = [], []\n",
        "    \n",
        "    #Training phase.    \n",
        "    for epoch in range(EPOCHS):\n",
        "\n",
        "        train_running_loss = 0\n",
        "        correct_train = 0\n",
        "        \n",
        "        model.train().cuda()\n",
        "        \n",
        "        for inputs, labels in train_loader:\n",
        "            inputs, labels = inputs.to(device), labels.to(device)\n",
        "            \n",
        "            \"\"\" for every mini-batch during the training phase, we typically want to explicitly set the gradients \n",
        "            to zero before starting to do backpropragation \"\"\"\n",
        "            optimizer.zero_grad()\n",
        "            \n",
        "            # Start the forward pass\n",
        "            prediction = model(inputs)\n",
        "                        \n",
        "            loss = criterion(prediction, labels)\n",
        "          \n",
        "            # do backpropagation and update weights with step()\n",
        "            loss.backward()         \n",
        "            optimizer.step()\n",
        "            \n",
        "            # print('outputs on which to apply torch.max ', prediction)\n",
        "            # find the maximum along the rows, use dim=1 to torch.max()\n",
        "            _, predicted_outputs = torch.max(prediction.data, 1)\n",
        "            \n",
        "            # Update the running corrects \n",
        "            correct_train += (predicted_outputs == labels).float().sum().item()\n",
        "            \n",
        "            ''' Compute batch loss\n",
        "            multiply each average batch loss with batch-length. \n",
        "            The batch-length is inputs.size(0) which gives the number total images in each batch. \n",
        "            Essentially I am un-averaging the previously calculated Loss '''\n",
        "            train_running_loss += (loss.data.item() * inputs.shape[0])\n",
        "\n",
        "\n",
        "        train_epoch_loss = train_running_loss / train_samples_num\n",
        "        \n",
        "        train_costs.append(train_epoch_loss)\n",
        "        \n",
        "        train_acc =  correct_train / train_samples_num\n",
        "\n",
        "        # Now check trained weights on the validation set\n",
        "        val_running_loss = 0\n",
        "        correct_val = 0\n",
        "      \n",
        "        model.eval().cuda()\n",
        "    \n",
        "        with torch.no_grad():\n",
        "            for inputs, labels in val_loader:\n",
        "                inputs, labels = inputs.to(device), labels.to(device)\n",
        "\n",
        "                # Forward pass.\n",
        "                prediction = model(inputs)\n",
        "\n",
        "                # Compute the loss.\n",
        "                loss = criterion(prediction, labels)\n",
        "\n",
        "                # Compute validation accuracy.\n",
        "                _, predicted_outputs = torch.max(prediction.data, 1)\n",
        "                correct_val += (predicted_outputs == labels).float().sum().item()\n",
        "\n",
        "            # Compute batch loss.\n",
        "            val_running_loss += (loss.data.item() * inputs.shape[0])\n",
        "\n",
        "            val_epoch_loss = val_running_loss / val_samples_num\n",
        "            val_costs.append(val_epoch_loss)\n",
        "            val_acc =  correct_val / val_samples_num\n",
        "        \n",
        "        info = \"[Epoch {}/{}]: train-loss = {:0.6f} | train-acc = {:0.3f} | val-loss = {:0.6f} | val-acc = {:0.3f}\"\n",
        "        \n",
        "        print(info.format(epoch+1, EPOCHS, train_epoch_loss, train_acc, val_epoch_loss, val_acc))\n",
        "        \n",
        "        torch.save(model.state_dict(), '/content/checkpoint_gpu_{}'.format(epoch + 1)) \n",
        "                                                                \n",
        "    torch.save(model.state_dict(), '/content/resnet-56_weights_gpu')  \n",
        "        \n",
        "    return train_costs, val_costs\n",
        "\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JJaEBgbt9INy",
        "outputId": "56270746-ac9e-4ffd-b481-b3d0b901ca82"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[Epoch 1/15]: train-loss = 0.876177 | train-acc = 0.692 | val-loss = 0.001271 | val-acc = 0.734\n",
            "[Epoch 2/15]: train-loss = 0.694989 | train-acc = 0.759 | val-loss = 0.002828 | val-acc = 0.769\n",
            "[Epoch 3/15]: train-loss = 0.580110 | train-acc = 0.800 | val-loss = 0.000333 | val-acc = 0.778\n",
            "[Epoch 4/15]: train-loss = 0.492105 | train-acc = 0.829 | val-loss = 0.001201 | val-acc = 0.780\n",
            "[Epoch 5/15]: train-loss = 0.416747 | train-acc = 0.854 | val-loss = 0.001592 | val-acc = 0.810\n",
            "[Epoch 6/15]: train-loss = 0.351747 | train-acc = 0.877 | val-loss = 0.000720 | val-acc = 0.784\n",
            "[Epoch 7/15]: train-loss = 0.297035 | train-acc = 0.896 | val-loss = 0.001334 | val-acc = 0.794\n",
            "[Epoch 8/15]: train-loss = 0.248202 | train-acc = 0.912 | val-loss = 0.000497 | val-acc = 0.823\n",
            "[Epoch 9/15]: train-loss = 0.206233 | train-acc = 0.927 | val-loss = 0.001534 | val-acc = 0.814\n",
            "[Epoch 10/15]: train-loss = 0.169824 | train-acc = 0.940 | val-loss = 0.000273 | val-acc = 0.815\n",
            "[Epoch 11/15]: train-loss = 0.146284 | train-acc = 0.948 | val-loss = 0.000277 | val-acc = 0.818\n",
            "[Epoch 12/15]: train-loss = 0.126156 | train-acc = 0.955 | val-loss = 0.002989 | val-acc = 0.822\n",
            "[Epoch 13/15]: train-loss = 0.111513 | train-acc = 0.961 | val-loss = 0.000481 | val-acc = 0.824\n",
            "[Epoch 14/15]: train-loss = 0.109781 | train-acc = 0.962 | val-loss = 0.000927 | val-acc = 0.824\n",
            "[Epoch 15/15]: train-loss = 0.091795 | train-acc = 0.968 | val-loss = 0.003517 | val-acc = 0.833\n"
          ]
        }
      ],
      "source": [
        "# !pwd\n",
        "train_costs, val_costs = train_model()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kg45TKgb8N0k",
        "outputId": "fb8457c3-aaf2-42d1-d456-454012473c58"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "#Restore the model.\n",
        "model = ResNet56()\n",
        "model.load_state_dict(torch.load('/content/resnet-56_weights_gpu'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AODk608HNFbv"
      },
      "source": [
        "## Test the trained model on Test dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MPbkor1g8Q3r",
        "outputId": "66352854-1346-4280-d198-c31310685fc7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Test accuracy: 0.8344\n"
          ]
        }
      ],
      "source": [
        "test_samples_num = 10000\n",
        "correct = 0 \n",
        "\n",
        "model.eval().cuda()\n",
        "\n",
        "with  torch.no_grad():\n",
        "    for inputs, labels in test_loader:\n",
        "        inputs, labels = inputs.to(device), labels.to(device)\n",
        "        # Make predictions.\n",
        "        prediction = model(inputs)\n",
        "\n",
        "        # Retrieve predictions indexes.\n",
        "        _, predicted_class = torch.max(prediction.data, 1)\n",
        "\n",
        "        # Compute number of correct predictions.\n",
        "        correct += (predicted_class == labels).float().sum().item()\n",
        "\n",
        "test_accuracy = correct / test_samples_num\n",
        "print('Test accuracy: {}'.format(test_accuracy))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "ResNet_MyWork.ipynb",
      "provenance": []
    },
    "interpreter": {
      "hash": "36cf16204b8548560b1c020c4e8fb5b57f0e4c58016f52f2d4be01e192833930"
    },
    "kernelspec": {
      "display_name": "Python 3.9.10 64-bit",
      "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.10"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}