src/architecture/mobilenet_v1/pytorch/mobilenet_v1_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dqtqtCz46iBb"
   },
   "source": [
    "# I) Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- The paper [MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications](https://arxiv.org/pdf/1704.04861.pdf) describes a very small, low latency and efficient network architecture called MobileNet for mobile and embedded vision applications.\n",
    "- To do so, they use depthwise separable convolutions and a set of 2 hyperparameters (width and resolution multiplier).\n",
    "\n",
    "## 1) Depthwise separable convolution\n",
    "\n",
    "- To understand what a depthwise separable convolution really is, let's compare it to a normal convolution between a 12x12x3 input and 256 kernels of size 5x5x3.\n",
    "- A depthwise separable convolution is divided into 2 parts:\n",
    "    - **Depthwise convolution**.\n",
    "    - **Pointwise convolution**.\n",
    " \n",
    "**<ins>Depthwise convolution:</ins>**\n",
    "\n",
    "- In a normal convolution, **all channels** of a kernel are used to produce a feature map.\n",
    "- In a depthwise convolution, **each channel** of a kernel is used to produce a feature map.\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "    <img src=\"https://raw.githubusercontent.com/valoxe/image-storage-1/master/research-paper-summary/mobilenet-v1/1.png\" width=\"70%\">\n",
    "    <figcaption> Figure: Normal convolution</figcaption>\n",
    "</div>\n",
    "<br>\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "    <img src=\"https://raw.githubusercontent.com/valoxe/image-storage-1/master/research-paper-summary/mobilenet-v1/2.png\" width=\"70%\">\n",
    "    <figcaption> Figure: Depthwise convolution</figcaption>\n",
    "</div>\n",
    "\n",
    "**<ins>Pointwise convolution:</ins>**\n",
    "\n",
    "\n",
    "- To increase the number of channels in our output image to 256:\n",
    "    - In a **normal convolution**, we just have to use **256 filters of size 5x5x3**.\n",
    "    - In a **pointwise convolution**, we just have to use **256 filters of size 1x1x3**.\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "    <img src=\"https://raw.githubusercontent.com/valoxe/image-storage-1/master/research-paper-summary/mobilenet-v1/3.png\" width=\"70%\">\n",
    "    <figcaption> Figure: Normal convolution</figcaption>\n",
    "</div>\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "    <img src=\"https://raw.githubusercontent.com/valoxe/image-storage-1/master/research-paper-summary/mobilenet-v1/4.png\" width=\"70%\">\n",
    "    <figcaption> Figure: Pointwise convolution</figcaption>\n",
    "</div>\n",
    "\n",
    "- **What's the main difference between a depthwise separable convolution and normal convolution ?**\n",
    "- The **main difference** is the **number of computations**. In our example:\n",
    "    - For a **normal convolution**, we have ((8x8x5x5)x3)x256 = **1,228,800** operations.\n",
    "    - For a **depthwise separable convolution**, we have 4800 + 49,152 = **53,952** operations:\n",
    "        - in a **depthwise convolution**, (8x8x5x5)x3 = 4800 operations.\n",
    "        - in a **pointwise convolution**, ((8x8x1x1)x3)x256 = 49,152 operations.\n",
    "        \n",
    "- We can clearly see that a depthwise separable convolution is **less expensive** than a normal convolution (~22.7% less computations).\n",
    "- The reason is, in a normal convolution, we are **transforming the image 256 times** whereas in a depthwise separable convolution, we transform the image **once** and then **expand it 256 times** along the channel axis.\n",
    "- The authors made a comparison between a MobileNet with depthwise seperable convolution and one with normal convolution on Imagenet. **Turns out, the accuracy only dropped by 1% but has less parameters and operations**.\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "    <img src=\"https://raw.githubusercontent.com/valoxe/image-storage-1/master/research-paper-summary/mobilenet-v1/5.png\">\n",
    "</div>\n",
    " \n",
    "## 2) Hyperparameters\n",
    "\n",
    "- They demonstrated how to build smaller and faster MobileNets using width multiplier ($\\alpha$) and resolution multiplier ($\\rho$) by trading off a reasonable amount of accuracy to reduce size and latency.\n",
    "- The **width multiplier** ($\\alpha$) (also known as \"**depth multiplier**\") with values $\\{1, 0.75, 0.5, 0.25\\}$, thins a network uniformly at each layer leading to a **reduction in computational cost and number of parameters**. \n",
    "    - For example, if the width multiplier is 1, the network starts off with 32 channels and ends up with 1024.\n",
    "    - Using a width multiplier of 0.5 will halve the number of channels used in each layer resulting in a reduction of number of computations by a factor of 4 and a number of learnable parameters by a factor 3 (see Table 6). Therefore, the new model is faster but less accurate than the full model. \n",
    "- The **resolution multiplier** ($\\rho$) with values $\\{224, 192, 160, 128\\}$ reduces the input size leading to a **reduction in computational cost**.\n",
    "\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "    <img src=\"https://raw.githubusercontent.com/valoxe/image-storage-1/master/research-paper-summary/mobilenet-v1/6.png\">\n",
    "</div>\n",
    "<br>\n",
    "\n",
    "## 3) Architecture\n",
    "\n",
    "Here is MobileNet-V1 architecture:\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "    <img src=\"https://raw.githubusercontent.com/valoxe/image-storage-1/master/research-paper-summary/mobilenet-v1/7.png\" width=\"70%\">\n",
    "</div>\n",
    "<br>\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "    <img src=\"https://raw.githubusercontent.com/valoxe/image-storage-1/master/research-paper-summary/mobilenet-v1/8.png\" width=\"70%\">\n",
    "</div>\n",
    "<br>\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "    <img src=\"https://raw.githubusercontent.com/valoxe/image-storage-1/master/research-paper-summary/mobilenet-v1/9.png\" width=\"70%\">\n",
    "</div>\n",
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bVOv0X256iCA"
   },
   "source": [
    "# II) Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ttli5CUy6iCM"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "from collections import OrderedDict\n",
    "from IPython.display import clear_output\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": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hI2wFbh77Rcc"
   },
   "source": [
    "## a) Loading dataset / Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "grJe4-fY7U9M"
   },
   "outputs": [],
   "source": [
    "def load_cifar():\n",
    "    transform = transforms.Compose([transforms.Resize((96,96)),\n",
    "                                    transforms.ToTensor(),\n",
    "                                    transforms.Normalize(mean=[0.5], std=[0.5])])\n",
    "            \n",
    "    train_dataset = datasets.CIFAR10('./data', train=True, download=True, transform=transform)\n",
    "    test_dataset = datasets.CIFAR10('./data', train=False, download=True, transform=transform)\n",
    "\n",
    "    #Clear downloading message.\n",
    "    clear_output()\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: {}\".format(train_dataset[0][0].numpy().shape), end = '\\n\\n')\n",
    "    print(\"Training Set:   {} samples\".format(len(train_dataset)))\n",
    "    print(\"Validation Set:   {} samples\".format(len(val_dataset)))\n",
    "    print(\"Test Set:       {} samples\".format(len(test_dataset)))\n",
    "    \n",
    "    if torch.cuda.is_available():\n",
    "        BATCH_SIZE = 256\n",
    "    else:\n",
    "        BATCH_SIZE = 32\n",
    "\n",
    "    # Create iterator.\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=BATCH_SIZE, shuffle=True)\n",
    "    \n",
    "    # Delete the data/ folder.\n",
    "    shutil.rmtree('./data')\n",
    "    \n",
    "    return train_loader, val_loader, test_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 109,
     "referenced_widgets": [
      "c25b46e6e9d546b8a5d86e20fc23c8f4",
      "15e5145430f94172a6c19169a84a9002",
      "2633ea046ab04399a50569f1613be6ca",
      "0fe12bcef00c4a05ba92ea294ae92e89",
      "4c707883f00f4b15992538807c0a753c",
      "5a10eea885b64398b9e749e317f24976",
      "742e265652c24096ae258a670a5ae1fe",
      "61d6184bf4204fb1bb7619d6ac05142b"
     ]
    },
    "colab_type": "code",
    "id": "A2h3Gaus7XvO",
    "outputId": "13b4d885-f33a-4786-bd70-409f41a329e1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image Shape: (3, 96, 96)\n",
      "\n",
      "Training Set:   45000 samples\n",
      "Validation Set:   5000 samples\n",
      "Test Set:       10000 samples\n"
     ]
    }
   ],
   "source": [
    "train_loader, val_loader, test_loader = load_cifar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3vFaWiye7ZW5"
   },
   "source": [
    "## b) Architecture build"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "AN7yixMN6iCj"
   },
   "outputs": [],
   "source": [
    "class DSConv(nn.Module):\n",
    "    \n",
    "    def __init__(self, f_3x3, f_1x1, stride=1, padding=0):\n",
    "        super(DSConv, self).__init__()\n",
    "        \n",
    "        self.feature = nn.Sequential(OrderedDict([\n",
    "            ('dconv', nn.Conv2d(f_3x3,\n",
    "                                f_3x3,\n",
    "                                kernel_size=3,\n",
    "                                groups=f_3x3,\n",
    "                                stride=stride,\n",
    "                                padding=padding,\n",
    "                                bias=False\n",
    "                                )),\n",
    "            ('bn1', nn.BatchNorm2d(f_3x3)),\n",
    "            ('act1', nn.ReLU()),\n",
    "            ('pconv', nn.Conv2d(f_3x3,\n",
    "                                f_1x1,\n",
    "                                kernel_size=1,\n",
    "                                bias=False)),\n",
    "            ('bn2', nn.BatchNorm2d(f_1x1)),\n",
    "            ('act2', nn.ReLU())\n",
    "        ]))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        out = self.feature(x)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tfuaSdWg6iC6"
   },
   "outputs": [],
   "source": [
    "class MobileNet(nn.Module):\n",
    "    \"\"\"\n",
    "        MobileNet-V1 architecture for CIFAR-10.\n",
    "    \"\"\"\n",
    "    def __init__(self, channels, width_multiplier=1.0, num_classes=1000):\n",
    "        super(MobileNet, self).__init__()\n",
    "        \n",
    "        channels = [int(elt * width_multiplier) for elt in channels]\n",
    "        \n",
    "        self.conv = nn.Sequential(OrderedDict([\n",
    "            ('conv', nn.Conv2d(3, channels[0], kernel_size=3,\n",
    "                               stride=2, padding=1, bias=False)),\n",
    "            ('bn', nn.BatchNorm2d(channels[0])),\n",
    "            ('act', nn.ReLU()) \n",
    "        ]))\n",
    "        \n",
    "        self.features = nn.Sequential(OrderedDict([\n",
    "            ('dsconv1', DSConv(channels[0], channels[1], 1, 1)),\n",
    "            ('dsconv2', DSConv(channels[1], channels[2], 2, 1)),\n",
    "            ('dsconv3', DSConv(channels[2], channels[2], 1, 1)),\n",
    "            ('dsconv4', DSConv(channels[2], channels[3], 2, 1)),\n",
    "            ('dsconv5', DSConv(channels[3], channels[3], 1, 1)),\n",
    "            ('dsconv6', DSConv(channels[3], channels[4], 2, 1)),\n",
    "            ('dsconv7_a', DSConv(channels[4], channels[4], 1, 1)),\n",
    "            ('dsconv7_b', DSConv(channels[4], channels[4], 1, 1)),\n",
    "            ('dsconv7_c', DSConv(channels[4], channels[4], 1, 1)),\n",
    "            ('dsconv7_d', DSConv(channels[4], channels[4], 1, 1)),\n",
    "            ('dsconv7_e', DSConv(channels[4], channels[4], 1, 1)),\n",
    "            ('dsconv8', DSConv(channels[4], channels[5], 2, 1)),\n",
    "            ('dsconv9', DSConv(channels[5], channels[5], 1, 1))\n",
    "        ]))\n",
    "        \n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1,1))\n",
    "        self.linear = nn.Linear(channels[5], num_classes)\n",
    "       \n",
    "        \n",
    "    def forward(self, x):\n",
    "        out = self.conv(x)\n",
    "        out = self.features(out)\n",
    "        out = self.avgpool(out)\n",
    "        out = torch.flatten(out, 1)\n",
    "        out = self.linear(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "D-gKqqDa6iDN",
    "outputId": "649c4630-30ef-430d-9866-5a8e563a219a",
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1           [-1, 32, 48, 48]             864\n",
      "       BatchNorm2d-2           [-1, 32, 48, 48]              64\n",
      "              ReLU-3           [-1, 32, 48, 48]               0\n",
      "            Conv2d-4           [-1, 32, 48, 48]             288\n",
      "       BatchNorm2d-5           [-1, 32, 48, 48]              64\n",
      "              ReLU-6           [-1, 32, 48, 48]               0\n",
      "            Conv2d-7           [-1, 64, 48, 48]           2,048\n",
      "       BatchNorm2d-8           [-1, 64, 48, 48]             128\n",
      "              ReLU-9           [-1, 64, 48, 48]               0\n",
      "           DSConv-10           [-1, 64, 48, 48]               0\n",
      "           Conv2d-11           [-1, 64, 24, 24]             576\n",
      "      BatchNorm2d-12           [-1, 64, 24, 24]             128\n",
      "             ReLU-13           [-1, 64, 24, 24]               0\n",
      "           Conv2d-14          [-1, 128, 24, 24]           8,192\n",
      "      BatchNorm2d-15          [-1, 128, 24, 24]             256\n",
      "             ReLU-16          [-1, 128, 24, 24]               0\n",
      "           DSConv-17          [-1, 128, 24, 24]               0\n",
      "           Conv2d-18          [-1, 128, 24, 24]           1,152\n",
      "      BatchNorm2d-19          [-1, 128, 24, 24]             256\n",
      "             ReLU-20          [-1, 128, 24, 24]               0\n",
      "           Conv2d-21          [-1, 128, 24, 24]          16,384\n",
      "      BatchNorm2d-22          [-1, 128, 24, 24]             256\n",
      "             ReLU-23          [-1, 128, 24, 24]               0\n",
      "           DSConv-24          [-1, 128, 24, 24]               0\n",
      "           Conv2d-25          [-1, 128, 12, 12]           1,152\n",
      "      BatchNorm2d-26          [-1, 128, 12, 12]             256\n",
      "             ReLU-27          [-1, 128, 12, 12]               0\n",
      "           Conv2d-28          [-1, 256, 12, 12]          32,768\n",
      "      BatchNorm2d-29          [-1, 256, 12, 12]             512\n",
      "             ReLU-30          [-1, 256, 12, 12]               0\n",
      "           DSConv-31          [-1, 256, 12, 12]               0\n",
      "           Conv2d-32          [-1, 256, 12, 12]           2,304\n",
      "      BatchNorm2d-33          [-1, 256, 12, 12]             512\n",
      "             ReLU-34          [-1, 256, 12, 12]               0\n",
      "           Conv2d-35          [-1, 256, 12, 12]          65,536\n",
      "      BatchNorm2d-36          [-1, 256, 12, 12]             512\n",
      "             ReLU-37          [-1, 256, 12, 12]               0\n",
      "           DSConv-38          [-1, 256, 12, 12]               0\n",
      "           Conv2d-39            [-1, 256, 6, 6]           2,304\n",
      "      BatchNorm2d-40            [-1, 256, 6, 6]             512\n",
      "             ReLU-41            [-1, 256, 6, 6]               0\n",
      "           Conv2d-42            [-1, 512, 6, 6]         131,072\n",
      "      BatchNorm2d-43            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-44            [-1, 512, 6, 6]               0\n",
      "           DSConv-45            [-1, 512, 6, 6]               0\n",
      "           Conv2d-46            [-1, 512, 6, 6]           4,608\n",
      "      BatchNorm2d-47            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-48            [-1, 512, 6, 6]               0\n",
      "           Conv2d-49            [-1, 512, 6, 6]         262,144\n",
      "      BatchNorm2d-50            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-51            [-1, 512, 6, 6]               0\n",
      "           DSConv-52            [-1, 512, 6, 6]               0\n",
      "           Conv2d-53            [-1, 512, 6, 6]           4,608\n",
      "      BatchNorm2d-54            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-55            [-1, 512, 6, 6]               0\n",
      "           Conv2d-56            [-1, 512, 6, 6]         262,144\n",
      "      BatchNorm2d-57            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-58            [-1, 512, 6, 6]               0\n",
      "           DSConv-59            [-1, 512, 6, 6]               0\n",
      "           Conv2d-60            [-1, 512, 6, 6]           4,608\n",
      "      BatchNorm2d-61            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-62            [-1, 512, 6, 6]               0\n",
      "           Conv2d-63            [-1, 512, 6, 6]         262,144\n",
      "      BatchNorm2d-64            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-65            [-1, 512, 6, 6]               0\n",
      "           DSConv-66            [-1, 512, 6, 6]               0\n",
      "           Conv2d-67            [-1, 512, 6, 6]           4,608\n",
      "      BatchNorm2d-68            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-69            [-1, 512, 6, 6]               0\n",
      "           Conv2d-70            [-1, 512, 6, 6]         262,144\n",
      "      BatchNorm2d-71            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-72            [-1, 512, 6, 6]               0\n",
      "           DSConv-73            [-1, 512, 6, 6]               0\n",
      "           Conv2d-74            [-1, 512, 6, 6]           4,608\n",
      "      BatchNorm2d-75            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-76            [-1, 512, 6, 6]               0\n",
      "           Conv2d-77            [-1, 512, 6, 6]         262,144\n",
      "      BatchNorm2d-78            [-1, 512, 6, 6]           1,024\n",
      "             ReLU-79            [-1, 512, 6, 6]               0\n",
      "           DSConv-80            [-1, 512, 6, 6]               0\n",
      "           Conv2d-81            [-1, 512, 3, 3]           4,608\n",
      "      BatchNorm2d-82            [-1, 512, 3, 3]           1,024\n",
      "             ReLU-83            [-1, 512, 3, 3]               0\n",
      "           Conv2d-84           [-1, 1024, 3, 3]         524,288\n",
      "      BatchNorm2d-85           [-1, 1024, 3, 3]           2,048\n",
      "             ReLU-86           [-1, 1024, 3, 3]               0\n",
      "           DSConv-87           [-1, 1024, 3, 3]               0\n",
      "           Conv2d-88           [-1, 1024, 3, 3]           9,216\n",
      "      BatchNorm2d-89           [-1, 1024, 3, 3]           2,048\n",
      "             ReLU-90           [-1, 1024, 3, 3]               0\n",
      "           Conv2d-91           [-1, 1024, 3, 3]       1,048,576\n",
      "      BatchNorm2d-92           [-1, 1024, 3, 3]           2,048\n",
      "             ReLU-93           [-1, 1024, 3, 3]               0\n",
      "           DSConv-94           [-1, 1024, 3, 3]               0\n",
      "AdaptiveAvgPool2d-95           [-1, 1024, 1, 1]               0\n",
      "           Linear-96                 [-1, 1000]       1,025,000\n",
      "================================================================\n",
      "Total params: 4,231,976\n",
      "Trainable params: 4,231,976\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.11\n",
      "Forward/backward pass size (MB): 25.01\n",
      "Params size (MB): 16.14\n",
      "Estimated Total Size (MB): 41.26\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "def MobileNetV1():\n",
    "    return MobileNet(channels=[32, 64, 128, 256, 512, 1024], width_multiplier=1)\n",
    "\n",
    "model = MobileNetV1()\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "model.to(device)\n",
    "summary(model, (3, 96, 96))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jjOKqUHg6iDh"
   },
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bXH173707pNf"
   },
   "source": [
    "## c) Training the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Wsptl8Kq7l_H"
   },
   "outputs": [],
   "source": [
    "def train_model():\n",
    "    EPOCHS = 15\n",
    "    nb_examples = 45000\n",
    "    nb_val_examples = 5000\n",
    "    train_costs, val_costs = [], []\n",
    "    \n",
    "    #Training phase.\n",
    "    \n",
    "    for epoch in range(EPOCHS):\n",
    "\n",
    "        train_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",
    "            # Zero the parameter gradients.\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            # Forward pass.\n",
    "            prediction = model(inputs)\n",
    "            \n",
    "            # Compute the loss.\n",
    "            loss = criterion(prediction, labels)\n",
    "          \n",
    "            # Backward pass.\n",
    "            loss.backward()\n",
    "            \n",
    "            # Optimize.\n",
    "            optimizer.step()\n",
    "            \n",
    "            # Compute training accuracy.\n",
    "            _, predicted = torch.max(prediction.data, 1)\n",
    "            correct_train += (predicted == labels).float().sum().item()\n",
    "            \n",
    "            # Compute batch loss.\n",
    "            train_loss += (loss.data.item() * inputs.shape[0])\n",
    "\n",
    "\n",
    "        train_loss /= nb_examples\n",
    "        train_costs.append(train_loss)\n",
    "        train_acc =  correct_train / nb_examples\n",
    "\n",
    "        val_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 training accuracy.\n",
    "                _, predicted = torch.max(prediction.data, 1)\n",
    "                correct_val += (predicted == labels).float().sum().item()\n",
    "\n",
    "            # Compute batch loss.\n",
    "            val_loss += (loss.data.item() * inputs.shape[0])\n",
    "\n",
    "            val_loss /= nb_val_examples\n",
    "            val_costs.append(val_loss)\n",
    "            val_acc =  correct_val / nb_val_examples\n",
    "        \n",
    "        info = \"[Epoch {}/{}]: train-loss = {:0.6f} | train-acc = {:0.3f} | val-loss = {:0.6f} | val-acc = {:0.3f}\"\n",
    "        print(info.format(epoch+1, EPOCHS, train_loss, train_acc, val_loss, val_acc))\n",
    "        torch.save(model.state_dict(), 'save_weights/checkpoint_gpu_{}'.format(epoch + 1)) \n",
    "                                                                \n",
    "    torch.save(model.state_dict(), 'save_weights/mobilenet-v1_weights_gpu')  \n",
    "        \n",
    "    return train_costs, val_costs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 293
    },
    "colab_type": "code",
    "id": "oWyNogJk7sdT",
    "outputId": "cc277d01-a9d1-46b8-c52b-ede9ee94c4c4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch 1/15]: train-loss = 1.559708 | train-acc = 0.433 | val-loss = 0.032007 | val-acc = 0.586\n",
      "[Epoch 2/15]: train-loss = 1.004967 | train-acc = 0.645 | val-loss = 0.025247 | val-acc = 0.701\n",
      "[Epoch 3/15]: train-loss = 0.747347 | train-acc = 0.742 | val-loss = 0.021409 | val-acc = 0.745\n",
      "[Epoch 4/15]: train-loss = 0.596251 | train-acc = 0.793 | val-loss = 0.022274 | val-acc = 0.776\n",
      "[Epoch 5/15]: train-loss = 0.493434 | train-acc = 0.832 | val-loss = 0.014880 | val-acc = 0.797\n",
      "[Epoch 6/15]: train-loss = 0.415764 | train-acc = 0.856 | val-loss = 0.012545 | val-acc = 0.808\n",
      "[Epoch 7/15]: train-loss = 0.356464 | train-acc = 0.877 | val-loss = 0.014014 | val-acc = 0.803\n",
      "[Epoch 8/15]: train-loss = 0.318513 | train-acc = 0.891 | val-loss = 0.011231 | val-acc = 0.808\n",
      "[Epoch 9/15]: train-loss = 0.265170 | train-acc = 0.909 | val-loss = 0.017637 | val-acc = 0.824\n",
      "[Epoch 10/15]: train-loss = 0.234187 | train-acc = 0.920 | val-loss = 0.021207 | val-acc = 0.827\n",
      "[Epoch 11/15]: train-loss = 0.203297 | train-acc = 0.930 | val-loss = 0.020546 | val-acc = 0.829\n",
      "[Epoch 12/15]: train-loss = 0.184567 | train-acc = 0.937 | val-loss = 0.015497 | val-acc = 0.827\n",
      "[Epoch 13/15]: train-loss = 0.167254 | train-acc = 0.942 | val-loss = 0.018095 | val-acc = 0.824\n",
      "[Epoch 14/15]: train-loss = 0.142094 | train-acc = 0.951 | val-loss = 0.019001 | val-acc = 0.826\n",
      "[Epoch 15/15]: train-loss = 0.128344 | train-acc = 0.955 | val-loss = 0.019820 | val-acc = 0.817\n"
     ]
    }
   ],
   "source": [
    "train_costs, val_costs = train_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "XHHz8ipq7uC2",
    "outputId": "97e024ef-b5f1-4f59-8776-90acc1d9ab4d"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 20,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Restore the model.\n",
    "model = MobileNetV1()\n",
    "model.load_state_dict(torch.load('save_weights/mobilenet-v1_weights_gpu'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "mEuZdnjA75uU",
    "outputId": "06c941a1-1545-404c-dd36-4152aeb67f1a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy: 0.8173\n"
     ]
    }
   ],
   "source": [
    "nb_test_examples = 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 / nb_test_examples\n",
    "print('Test accuracy: {}'.format(test_accuracy))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "mobilenet_v1_pytorch.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.6.9"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0fe12bcef00c4a05ba92ea294ae92e89": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_61d6184bf4204fb1bb7619d6ac05142b",
      "placeholder": "​",
      "style": "IPY_MODEL_742e265652c24096ae258a670a5ae1fe",
      "value": " 170500096/? [00:21&lt;00:00, 52479218.92it/s]"
     }
    },
    "15e5145430f94172a6c19169a84a9002": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "2633ea046ab04399a50569f1613be6ca": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "IntProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "IntProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "info",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5a10eea885b64398b9e749e317f24976",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_4c707883f00f4b15992538807c0a753c",
      "value": 1
     }
    },
    "4c707883f00f4b15992538807c0a753c": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": "initial"
     }
    },
    "5a10eea885b64398b9e749e317f24976": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "61d6184bf4204fb1bb7619d6ac05142b": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "742e265652c24096ae258a670a5ae1fe": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "c25b46e6e9d546b8a5d86e20fc23c8f4": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_2633ea046ab04399a50569f1613be6ca",
       "IPY_MODEL_0fe12bcef00c4a05ba92ea294ae92e89"
      ],
      "layout": "IPY_MODEL_15e5145430f94172a6c19169a84a9002"
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}