validated/vision/classification/resnet/train_resnet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Notebook for ResNet v1 and v2 models on ImageNet Dataset\n",
    "\n",
    "## Overview\n",
    "Use this notebook to train a ResNet model from scratch. **Make sure to have the ImageNet dataset prepared** according to the guidelines in the dataset section in [ResNet readme](README.md) before proceeding."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "The following dependencies need to be installed before proceeding.\n",
    "* mxnet - `pip install mxnet-cu90mkl` (tested on this version GPU, can use other versions)\n",
    "* gluoncv - `pip install gluoncv`\n",
    "* numpy - `pip install numpy`\n",
    "* matplotlib - `pip install matplotlib`\n",
    "\n",
    "In order to train the model with a python script: \n",
    "* Generate the script : In Jupyter Notebook browser, go to File -> Download as -> Python (.py)\n",
    "* Run the script: `python train_resnet.py`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import dependencies\n",
    "Verify that all dependencies are installed using the cell below. Continue if no errors encountered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "\n",
    "import argparse, time, logging\n",
    "\n",
    "import mxnet as mx\n",
    "import numpy as np\n",
    "from mxnet import gluon, nd\n",
    "from mxnet import autograd as ag\n",
    "from mxnet.gluon import nn\n",
    "from mxnet.gluon.data.vision import transforms\n",
    "\n",
    "from gluoncv.data import imagenet\n",
    "from gluoncv.utils import makedirs, TrainingHistory\n",
    "\n",
    "import os\n",
    "from mxnet.context import cpu\n",
    "from mxnet.gluon.block import HybridBlock\n",
    "from mxnet.gluon.contrib.nn import HybridConcurrent\n",
    "import multiprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Specify model, hyperparameters and save locations\n",
    "\n",
    "The training was done on a p3.16xlarge ec2 instance on AWS. It has 8 Nvidia Tesla V100 GPUs (16GB each) and Intel Xeon E5-2686 v4 @ 2.70GHz with 64 threads.\n",
    "\n",
    "The `batch_size` set below is per device. For multiple GPUs there are different batches in each GPU of size `batch_size` simultaneously.\n",
    "\n",
    "The rest of the parameters can be tuned to fit the needs of a user. The values shown below were used to train the model in the model zoo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify model - choose from (resnet18_v1,resnet18_v2,resnet34_v1,resnet34_v2,resnet50_v1,resnet50_v2,\n",
    "#resnet101_v1,resnet101_v2,resnet152_v1,resnet152_v2)\n",
    "model_name = 'resnet18_v1' \n",
    "\n",
    "# path to training and validation images to use\n",
    "data_dir = '/home/ubuntu/imagenet/img_dataset'\n",
    "\n",
    "# training batch size per device (CPU/GPU)\n",
    "# Used batch size = 64 for resnet18_v1,resnet18_v2,resnet34_v1,resnet34_v2,resnet50_v1,resnet50_v2,resnet101_v1,\n",
    "#resnet101_v2\n",
    "#Used batch size=32 for resnet152_v1,resnet152_v2\n",
    "batch_size = 64\n",
    "\n",
    "# number of GPUs to use (automatically detect the number of GPUs)\n",
    "num_gpus = len(mx.test_utils.list_gpus())\n",
    "\n",
    "# number of pre-processing workers (automatically detect the number of workers)\n",
    "num_workers = multiprocessing.cpu_count()\n",
    "\n",
    "# number of training epochs \n",
    "#used as 150 for all of the models , used 1 over here to show demo for 1 epoch\n",
    "num_epochs = 1\n",
    "\n",
    "# learning rate\n",
    "lr = 0.01\n",
    "\n",
    "# momentum value for optimizer\n",
    "momentum = 0.9\n",
    "\n",
    "# weight decay rate\n",
    "wd = 0.0002\n",
    "\n",
    "# decay rate of learning rate\n",
    "lr_decay = 0.1\n",
    "\n",
    "# interval for periodic learning rate decays\n",
    "lr_decay_period = 0\n",
    "\n",
    "# epoches at which learning rate decays\n",
    "lr_decay_epoch = '30,60,90'\n",
    "\n",
    "# mode in which to train the model. options are symbolic, imperative, hybrid\n",
    "mode = 'hybrid'\n",
    "\n",
    "# Number of batches to wait before logging\n",
    "log_interval = 50\n",
    "\n",
    "# frequency of model saving\n",
    "save_frequency = 10\n",
    "\n",
    "# directory of saved models\n",
    "save_dir = 'params'\n",
    "\n",
    "#directory of training logs\n",
    "logging_dir = 'logs'\n",
    "\n",
    "# the path to save the history plot\n",
    "save_plot_dir = '.'\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model definition in Gluon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "##This block contains definition for Resnet v1 and v2\n",
    "\n",
    "#Helpers\n",
    "def _conv3x3(channels, stride, in_channels):\n",
    "    return nn.Conv2D(channels, kernel_size=3, strides=stride, padding=1,\n",
    "                     use_bias=False, in_channels=in_channels)\n",
    "\n",
    "\n",
    "# Blocks\n",
    "class BasicBlockV1(HybridBlock):\n",
    "    r\"\"\"BasicBlock V1 from `\"Deep Residual Learning for Image Recognition\"\n",
    "    <http://arxiv.org/abs/1512.03385>`_ paper.\n",
    "    This is used for ResNet V1 for 18, 34 layers.\n",
    "    Parameters\n",
    "    ----------\n",
    "    channels : int\n",
    "        Number of output channels.\n",
    "    stride : int\n",
    "        Stride size.\n",
    "    downsample : bool, default False\n",
    "        Whether to downsample the input.\n",
    "    in_channels : int, default 0\n",
    "        Number of input channels. Default is 0, to infer from the graph.\n",
    "    \"\"\"\n",
    "    def __init__(self, channels, stride, downsample=False, in_channels=0, **kwargs):\n",
    "        super(BasicBlockV1, self).__init__(**kwargs)\n",
    "        self.body = nn.HybridSequential(prefix='')\n",
    "        self.body.add(_conv3x3(channels, stride, in_channels))\n",
    "        self.body.add(nn.BatchNorm())\n",
    "        self.body.add(nn.Activation('relu'))\n",
    "        self.body.add(_conv3x3(channels, 1, channels))\n",
    "        self.body.add(nn.BatchNorm())\n",
    "        if downsample:\n",
    "            self.downsample = nn.HybridSequential(prefix='')\n",
    "            self.downsample.add(nn.Conv2D(channels, kernel_size=1, strides=stride,\n",
    "                                          use_bias=False, in_channels=in_channels))\n",
    "            self.downsample.add(nn.BatchNorm())\n",
    "        else:\n",
    "            self.downsample = None\n",
    "\n",
    "    def hybrid_forward(self, F, x):\n",
    "        residual = x\n",
    "\n",
    "        x = self.body(x)\n",
    "\n",
    "        if self.downsample:\n",
    "            residual = self.downsample(residual)\n",
    "\n",
    "        x = F.Activation(residual+x, act_type='relu')\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class BottleneckV1(HybridBlock):\n",
    "    r\"\"\"Bottleneck V1 from `\"Deep Residual Learning for Image Recognition\"\n",
    "    <http://arxiv.org/abs/1512.03385>`_ paper.\n",
    "    This is used for ResNet V1 for 50, 101, 152 layers.\n",
    "    Parameters\n",
    "    ----------\n",
    "    channels : int\n",
    "        Number of output channels.\n",
    "    stride : int\n",
    "        Stride size.\n",
    "    downsample : bool, default False\n",
    "        Whether to downsample the input.\n",
    "    in_channels : int, default 0\n",
    "        Number of input channels. Default is 0, to infer from the graph.\n",
    "    \"\"\"\n",
    "    def __init__(self, channels, stride, downsample=False, in_channels=0, **kwargs):\n",
    "        super(BottleneckV1, self).__init__(**kwargs)\n",
    "        self.body = nn.HybridSequential(prefix='')\n",
    "        self.body.add(nn.Conv2D(channels//4, kernel_size=1, strides=stride))\n",
    "        self.body.add(nn.BatchNorm())\n",
    "        self.body.add(nn.Activation('relu'))\n",
    "        self.body.add(_conv3x3(channels//4, 1, channels//4))\n",
    "        self.body.add(nn.BatchNorm())\n",
    "        self.body.add(nn.Activation('relu'))\n",
    "        self.body.add(nn.Conv2D(channels, kernel_size=1, strides=1))\n",
    "        self.body.add(nn.BatchNorm())\n",
    "        if downsample:\n",
    "            self.downsample = nn.HybridSequential(prefix='')\n",
    "            self.downsample.add(nn.Conv2D(channels, kernel_size=1, strides=stride,\n",
    "                                          use_bias=False, in_channels=in_channels))\n",
    "            self.downsample.add(nn.BatchNorm())\n",
    "        else:\n",
    "            self.downsample = None\n",
    "\n",
    "    def hybrid_forward(self, F, x):\n",
    "        residual = x\n",
    "\n",
    "        x = self.body(x)\n",
    "\n",
    "        if self.downsample:\n",
    "            residual = self.downsample(residual)\n",
    "\n",
    "        x = F.Activation(x + residual, act_type='relu')\n",
    "        return x\n",
    "\n",
    "\n",
    "class BasicBlockV2(HybridBlock):\n",
    "    r\"\"\"BasicBlock V2 from\n",
    "    `\"Identity Mappings in Deep Residual Networks\"\n",
    "    <https://arxiv.org/abs/1603.05027>`_ paper.\n",
    "    This is used for ResNet V2 for 18, 34 layers.\n",
    "    Parameters\n",
    "    ----------\n",
    "    channels : int\n",
    "        Number of output channels.\n",
    "    stride : int\n",
    "        Stride size.\n",
    "    downsample : bool, default False\n",
    "        Whether to downsample the input.\n",
    "    in_channels : int, default 0\n",
    "        Number of input channels. Default is 0, to infer from the graph.\n",
    "    \"\"\"\n",
    "    def __init__(self, channels, stride, downsample=False, in_channels=0, **kwargs):\n",
    "        super(BasicBlockV2, self).__init__(**kwargs)\n",
    "        self.bn1 = nn.BatchNorm()\n",
    "        self.conv1 = _conv3x3(channels, stride, in_channels)\n",
    "        self.bn2 = nn.BatchNorm()\n",
    "        self.conv2 = _conv3x3(channels, 1, channels)\n",
    "        if downsample:\n",
    "            self.downsample = nn.Conv2D(channels, 1, stride, use_bias=False,\n",
    "                                        in_channels=in_channels)\n",
    "        else:\n",
    "            self.downsample = None\n",
    "\n",
    "    def hybrid_forward(self, F, x):\n",
    "        residual = x\n",
    "        x = self.bn1(x)\n",
    "        x = F.Activation(x, act_type='relu')\n",
    "        if self.downsample:\n",
    "            residual = self.downsample(x)\n",
    "        x = self.conv1(x)\n",
    "\n",
    "        x = self.bn2(x)\n",
    "        x = F.Activation(x, act_type='relu')\n",
    "        x = self.conv2(x)\n",
    "\n",
    "        return x + residual\n",
    "\n",
    "\n",
    "class BottleneckV2(HybridBlock):\n",
    "    r\"\"\"Bottleneck V2 from\n",
    "    `\"Identity Mappings in Deep Residual Networks\"\n",
    "    <https://arxiv.org/abs/1603.05027>`_ paper.\n",
    "    This is used for ResNet V2 for 50, 101, 152 layers.\n",
    "    Parameters\n",
    "    ----------\n",
    "    channels : int\n",
    "        Number of output channels.\n",
    "    stride : int\n",
    "        Stride size.\n",
    "    downsample : bool, default False\n",
    "        Whether to downsample the input.\n",
    "    in_channels : int, default 0\n",
    "        Number of input channels. Default is 0, to infer from the graph.\n",
    "    \"\"\"\n",
    "    def __init__(self, channels, stride, downsample=False, in_channels=0, **kwargs):\n",
    "        super(BottleneckV2, self).__init__(**kwargs)\n",
    "        self.bn1 = nn.BatchNorm()\n",
    "        self.conv1 = nn.Conv2D(channels//4, kernel_size=1, strides=1, use_bias=False)\n",
    "        self.bn2 = nn.BatchNorm()\n",
    "        self.conv2 = _conv3x3(channels//4, stride, channels//4)\n",
    "        self.bn3 = nn.BatchNorm()\n",
    "        self.conv3 = nn.Conv2D(channels, kernel_size=1, strides=1, use_bias=False)\n",
    "        if downsample:\n",
    "            self.downsample = nn.Conv2D(channels, 1, stride, use_bias=False,\n",
    "                                        in_channels=in_channels)\n",
    "        else:\n",
    "            self.downsample = None\n",
    "\n",
    "    def hybrid_forward(self, F, x):\n",
    "        residual = x\n",
    "        x = self.bn1(x)\n",
    "        x = F.Activation(x, act_type='relu')\n",
    "        if self.downsample:\n",
    "            residual = self.downsample(x)\n",
    "        x = self.conv1(x)\n",
    "\n",
    "        x = self.bn2(x)\n",
    "        x = F.Activation(x, act_type='relu')\n",
    "        x = self.conv2(x)\n",
    "\n",
    "        x = self.bn3(x)\n",
    "        x = F.Activation(x, act_type='relu')\n",
    "        x = self.conv3(x)\n",
    "\n",
    "        return x + residual\n",
    "\n",
    "\n",
    "# Nets\n",
    "class ResNetV1(HybridBlock):\n",
    "    r\"\"\"ResNet V1 model from\n",
    "    `\"Deep Residual Learning for Image Recognition\"\n",
    "    <http://arxiv.org/abs/1512.03385>`_ paper.\n",
    "    Parameters\n",
    "    ----------\n",
    "    block : HybridBlock\n",
    "        Class for the residual block. Options are BasicBlockV1, BottleneckV1.\n",
    "    layers : list of int\n",
    "        Numbers of layers in each block\n",
    "    channels : list of int\n",
    "        Numbers of channels in each block. Length should be one larger than layers list.\n",
    "    classes : int, default 1000\n",
    "        Number of classification classes.\n",
    "    thumbnail : bool, default False\n",
    "        Enable thumbnail.\n",
    "    \"\"\"\n",
    "    def __init__(self, block, layers, channels, classes=1000, thumbnail=False, **kwargs):\n",
    "        super(ResNetV1, self).__init__(**kwargs)\n",
    "        assert len(layers) == len(channels) - 1\n",
    "        with self.name_scope():\n",
    "            self.features = nn.HybridSequential(prefix='')\n",
    "            if thumbnail:\n",
    "                self.features.add(_conv3x3(channels[0], 1, 0))\n",
    "            else:\n",
    "                self.features.add(nn.Conv2D(channels[0], 7, 2, 3, use_bias=False))\n",
    "                self.features.add(nn.BatchNorm())\n",
    "                self.features.add(nn.Activation('relu'))\n",
    "                self.features.add(nn.MaxPool2D(3, 2, 1))\n",
    "\n",
    "            for i, num_layer in enumerate(layers):\n",
    "                stride = 1 if i == 0 else 2\n",
    "                self.features.add(self._make_layer(block, num_layer, channels[i+1],\n",
    "                                                   stride, i+1, in_channels=channels[i]))\n",
    "            self.features.add(nn.GlobalAvgPool2D())\n",
    "\n",
    "            self.output = nn.Dense(classes, in_units=channels[-1])\n",
    "\n",
    "    def _make_layer(self, block, layers, channels, stride, stage_index, in_channels=0):\n",
    "        layer = nn.HybridSequential(prefix='stage%d_'%stage_index)\n",
    "        with layer.name_scope():\n",
    "            layer.add(block(channels, stride, channels != in_channels, in_channels=in_channels,\n",
    "                            prefix=''))\n",
    "            for _ in range(layers-1):\n",
    "                layer.add(block(channels, 1, False, in_channels=channels, prefix=''))\n",
    "        return layer\n",
    "\n",
    "    def hybrid_forward(self, F, x):\n",
    "        x = self.features(x)\n",
    "        x = self.output(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class ResNetV2(HybridBlock):\n",
    "    r\"\"\"ResNet V2 model from\n",
    "    `\"Identity Mappings in Deep Residual Networks\"\n",
    "    <https://arxiv.org/abs/1603.05027>`_ paper.\n",
    "    Parameters\n",
    "    ----------\n",
    "    block : HybridBlock\n",
    "        Class for the residual block. Options are BasicBlockV1, BottleneckV1.\n",
    "    layers : list of int\n",
    "        Numbers of layers in each block\n",
    "    channels : list of int\n",
    "        Numbers of channels in each block. Length should be one larger than layers list.\n",
    "    classes : int, default 1000\n",
    "        Number of classification classes.\n",
    "    thumbnail : bool, default False\n",
    "        Enable thumbnail.\n",
    "    \"\"\"\n",
    "    def __init__(self, block, layers, channels, classes=1000, thumbnail=False, **kwargs):\n",
    "        super(ResNetV2, self).__init__(**kwargs)\n",
    "        assert len(layers) == len(channels) - 1\n",
    "        with self.name_scope():\n",
    "            self.features = nn.HybridSequential(prefix='')\n",
    "            self.features.add(nn.BatchNorm(scale=False, center=False))\n",
    "            if thumbnail:\n",
    "                self.features.add(_conv3x3(channels[0], 1, 0))\n",
    "            else:\n",
    "                self.features.add(nn.Conv2D(channels[0], 7, 2, 3, use_bias=False))\n",
    "                self.features.add(nn.BatchNorm())\n",
    "                self.features.add(nn.Activation('relu'))\n",
    "                self.features.add(nn.MaxPool2D(3, 2, 1))\n",
    "\n",
    "            in_channels = channels[0]\n",
    "            for i, num_layer in enumerate(layers):\n",
    "                stride = 1 if i == 0 else 2\n",
    "                self.features.add(self._make_layer(block, num_layer, channels[i+1],\n",
    "                                                   stride, i+1, in_channels=in_channels))\n",
    "                in_channels = channels[i+1]\n",
    "            self.features.add(nn.BatchNorm())\n",
    "            self.features.add(nn.Activation('relu'))\n",
    "            self.features.add(nn.GlobalAvgPool2D())\n",
    "            self.features.add(nn.Flatten())\n",
    "\n",
    "            self.output = nn.Dense(classes, in_units=in_channels)\n",
    "\n",
    "    def _make_layer(self, block, layers, channels, stride, stage_index, in_channels=0):\n",
    "        layer = nn.HybridSequential(prefix='stage%d_'%stage_index)\n",
    "        with layer.name_scope():\n",
    "            layer.add(block(channels, stride, channels != in_channels, in_channels=in_channels,\n",
    "                            prefix=''))\n",
    "            for _ in range(layers-1):\n",
    "                layer.add(block(channels, 1, False, in_channels=channels, prefix=''))\n",
    "        return layer\n",
    "\n",
    "    def hybrid_forward(self, F, x):\n",
    "        x = self.features(x)\n",
    "        x = self.output(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# Specification\n",
    "resnet_spec = {18: ('basic_block', [2, 2, 2, 2], [64, 64, 128, 256, 512]),\n",
    "               34: ('basic_block', [3, 4, 6, 3], [64, 64, 128, 256, 512]),\n",
    "               50: ('bottle_neck', [3, 4, 6, 3], [64, 256, 512, 1024, 2048]),\n",
    "               101: ('bottle_neck', [3, 4, 23, 3], [64, 256, 512, 1024, 2048]),\n",
    "               152: ('bottle_neck', [3, 8, 36, 3], [64, 256, 512, 1024, 2048])}\n",
    "\n",
    "resnet_net_versions = [ResNetV1, ResNetV2]\n",
    "resnet_block_versions = [{'basic_block': BasicBlockV1, 'bottle_neck': BottleneckV1},\n",
    "                         {'basic_block': BasicBlockV2, 'bottle_neck': BottleneckV2}]\n",
    "\n",
    "\n",
    "# Constructor\n",
    "def get_resnet(version, num_layers, **kwargs):\n",
    "    r\"\"\"ResNet V1 model from `\"Deep Residual Learning for Image Recognition\"\n",
    "    <http://arxiv.org/abs/1512.03385>`_ paper.\n",
    "    ResNet V2 model from `\"Identity Mappings in Deep Residual Networks\"\n",
    "    <https://arxiv.org/abs/1603.05027>`_ paper.\n",
    "    Parameters\n",
    "    ----------\n",
    "    version : int\n",
    "        Version of ResNet. Options are 1, 2.\n",
    "    num_layers : int\n",
    "        Numbers of layers. Options are 18, 34, 50, 101, 152.\n",
    "    \"\"\"\n",
    "    assert num_layers in resnet_spec, \\\n",
    "        \"Invalid number of layers: %d. Options are %s\"%(\n",
    "            num_layers, str(resnet_spec.keys()))\n",
    "    block_type, layers, channels = resnet_spec[num_layers]\n",
    "    assert version >= 1 and version <= 2, \\\n",
    "        \"Invalid resnet version: %d. Options are 1 and 2.\"%version\n",
    "    resnet_class = resnet_net_versions[version-1]\n",
    "    block_class = resnet_block_versions[version-1][block_type]\n",
    "    net = resnet_class(block_class, layers, channels, **kwargs)\n",
    "    return net\n",
    "\n",
    "def resnet18_v1(**kwargs):\n",
    "    r\"\"\"ResNet-18 V1 model from `\"Deep Residual Learning for Image Recognition\"\n",
    "    <http://arxiv.org/abs/1512.03385>`_ paper.\n",
    "    \"\"\"\n",
    "    return get_resnet(1, 18, **kwargs)\n",
    "\n",
    "def resnet34_v1(**kwargs):\n",
    "    r\"\"\"ResNet-34 V1 model from `\"Deep Residual Learning for Image Recognition\"\n",
    "    <http://arxiv.org/abs/1512.03385>`_ paper.\n",
    "    \"\"\"\n",
    "    return get_resnet(1, 34, **kwargs)\n",
    "\n",
    "def resnet50_v1(**kwargs):\n",
    "    r\"\"\"ResNet-50 V1 model from `\"Deep Residual Learning for Image Recognition\"\n",
    "    <http://arxiv.org/abs/1512.03385>`_ paper.\n",
    "    \"\"\"\n",
    "    return get_resnet(1, 50, **kwargs)\n",
    "\n",
    "def resnet101_v1(**kwargs):\n",
    "    r\"\"\"ResNet-101 V1 model from `\"Deep Residual Learning for Image Recognition\"\n",
    "    <http://arxiv.org/abs/1512.03385>`_ paper.\n",
    "    \"\"\"\n",
    "    return get_resnet(1, 101, **kwargs)\n",
    "\n",
    "def resnet152_v1(**kwargs):\n",
    "    r\"\"\"ResNet-152 V1 model from `\"Deep Residual Learning for Image Recognition\"\n",
    "    <http://arxiv.org/abs/1512.03385>`_ paper.\n",
    "    \"\"\"\n",
    "    return get_resnet(1, 152, **kwargs)\n",
    "\n",
    "def resnet18_v2(**kwargs):\n",
    "    r\"\"\"ResNet-18 V2 model from `\"Identity Mappings in Deep Residual Networks\"\n",
    "    <https://arxiv.org/abs/1603.05027>`_ paper.\n",
    "    \"\"\"\n",
    "    return get_resnet(2, 18, **kwargs)\n",
    "\n",
    "def resnet34_v2(**kwargs):\n",
    "    r\"\"\"ResNet-34 V2 model from `\"Identity Mappings in Deep Residual Networks\"\n",
    "    <https://arxiv.org/abs/1603.05027>`_ paper.\n",
    "    \"\"\"\n",
    "    return get_resnet(2, 34, **kwargs)\n",
    "\n",
    "def resnet50_v2(**kwargs):\n",
    "    r\"\"\"ResNet-50 V2 model from `\"Identity Mappings in Deep Residual Networks\"\n",
    "    <https://arxiv.org/abs/1603.05027>`_ paper.\n",
    "    \"\"\"\n",
    "    return get_resnet(2, 50, **kwargs)\n",
    "\n",
    "def resnet101_v2(**kwargs):\n",
    "    r\"\"\"ResNet-101 V2 model from `\"Identity Mappings in Deep Residual Networks\"\n",
    "    <https://arxiv.org/abs/1603.05027>`_ paper.\n",
    "    \"\"\"\n",
    "    return get_resnet(2, 101, **kwargs)\n",
    "\n",
    "def resnet152_v2(**kwargs):\n",
    "    r\"\"\"ResNet-152 V2 model from `\"Identity Mappings in Deep Residual Networks\"\n",
    "    <https://arxiv.org/abs/1603.05027>`_ paper.\n",
    "    \"\"\"\n",
    "    return get_resnet(2, 152, **kwargs)\n",
    "\n",
    "models = {    'resnet18_v1': resnet18_v1,\n",
    "              'resnet34_v1': resnet34_v1,\n",
    "              'resnet50_v1': resnet50_v1,\n",
    "              'resnet101_v1': resnet101_v1,\n",
    "              'resnet152_v1': resnet152_v1,\n",
    "              'resnet18_v2': resnet18_v2,\n",
    "              'resnet34_v2': resnet34_v2,\n",
    "              'resnet50_v2': resnet50_v2,\n",
    "              'resnet101_v2': resnet101_v2,\n",
    "              'resnet152_v2': resnet152_v2\n",
    "         }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper code\n",
    "Define context, optimizer, accuracy metrics, retireve gluon model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify logging function\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "\n",
    "# Specify classes (1000 for ImageNet)\n",
    "classes = 1000\n",
    "# Extrapolate batches to all devices\n",
    "batch_size *= max(1, num_gpus)\n",
    "# Define context\n",
    "context = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu()]\n",
    "\n",
    "lr_decay_epoch = [int(i) for i in lr_decay_epoch.split(',')] + [np.inf]\n",
    "\n",
    "kwargs = { 'classes': classes}\n",
    "\n",
    "# Define optimizer (nag = Nestrov Accelerated Gradient)\n",
    "optimizer = 'nag'\n",
    "optimizer_params = {'learning_rate': lr, 'wd': wd, 'momentum': momentum}\n",
    "kwargs['thumbnail'] = False\n",
    "\n",
    "# Retrieve gluon model\n",
    "net = models[model_name](**kwargs)\n",
    "\n",
    "# Define accuracy measures - top1 error and top5 error\n",
    "acc_top1 = mx.metric.Accuracy()\n",
    "acc_top5 = mx.metric.TopKAccuracy(5)\n",
    "train_history = TrainingHistory(['training-top1-err', 'training-top5-err',\n",
    "                                 'validation-top1-err', 'validation-top5-err'])\n",
    "makedirs(save_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define preprocessing functions\n",
    "`preprocess_train_data(normalize, jitter_param, lighting_param)` : Do pre-processing and data augmentation of train images -> take random crops of size 224x224, do random left right flips, jitter image color and lighting, mormalize image\n",
    "\n",
    "`preprocess_test_data(normalize)` : Pre-process validation images -> resize to size 256x256, take center crop of size 224x224, normalize image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "jitter_param = 0.4\n",
    "lighting_param = 0.1\n",
    "\n",
    "# Input pre-processing for train data\n",
    "def preprocess_train_data(normalize, jitter_param, lighting_param):\n",
    "    transform_train = transforms.Compose([\n",
    "        transforms.Resize(480),\n",
    "        transforms.RandomResizedCrop(224),\n",
    "        transforms.RandomFlipLeftRight(),\n",
    "        transforms.RandomColorJitter(brightness=jitter_param, contrast=jitter_param,\n",
    "                                     saturation=jitter_param),\n",
    "        transforms.RandomLighting(lighting_param),\n",
    "        transforms.ToTensor(),\n",
    "        normalize\n",
    "    ])\n",
    "    return transform_train\n",
    "\n",
    "# Input pre-processing for validation data\n",
    "def preprocess_test_data(normalize):\n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        normalize\n",
    "    ])\n",
    "    return transform_test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define test function\n",
    "`test(ctx, val_data)` : Computes and returns validation errors on `val_data` using `ctx` context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test function\n",
    "def test(ctx, val_data):\n",
    "    # Reset accuracy metrics\n",
    "    acc_top1.reset()\n",
    "    acc_top5.reset()\n",
    "    for i, batch in enumerate(val_data):\n",
    "        # Load validation batch\n",
    "        data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)\n",
    "        label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0)\n",
    "        # Perform forward pass\n",
    "        outputs = [net(X) for X in data]\n",
    "        # Update accuracy metrics\n",
    "        acc_top1.update(label, outputs)\n",
    "        acc_top5.update(label, outputs)\n",
    "    # Retrieve and return top1 and top5 errors\n",
    "    _, top1 = acc_top1.get()\n",
    "    _, top5 = acc_top5.get()\n",
    "    return (1-top1, 1-top5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define train function\n",
    "`train(epochs, ctx)` : Train model for `epochs` epochs using `ctx` context, log training progress, compute and display validation errors after each epoch, take periodic snapshots of the model, generates training plot "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epochs, ctx):\n",
    "    if isinstance(ctx, mx.Context):\n",
    "        ctx = [ctx]\n",
    "    # Initialize network\n",
    "    net.initialize(mx.init.Xavier(magnitude=2), ctx=ctx)\n",
    "    # Prepare train and validation batches\n",
    "    transform_train = preprocess_train_data(normalize, jitter_param, lighting_param)\n",
    "    transform_test = preprocess_test_data(normalize)\n",
    "    train_data = gluon.data.DataLoader(\n",
    "        imagenet.classification.ImageNet(data_dir, train=True).transform_first(transform_train),\n",
    "        batch_size=batch_size, shuffle=True, last_batch='discard', num_workers=num_workers)\n",
    "    val_data = gluon.data.DataLoader(\n",
    "        imagenet.classification.ImageNet(data_dir, train=False).transform_first(transform_test),\n",
    "        batch_size=batch_size, shuffle=False, num_workers=num_workers)\n",
    "\n",
    "    # Define trainer\n",
    "    trainer = gluon.Trainer(net.collect_params(), optimizer, optimizer_params)\n",
    "    # Define loss\n",
    "    L = gluon.loss.SoftmaxCrossEntropyLoss()\n",
    "\n",
    "    lr_decay_count = 0\n",
    "\n",
    "    best_val_score = 1\n",
    "    # Main training loop - loop over epochs\n",
    "    for epoch in range(epochs):\n",
    "        tic = time.time()\n",
    "        # Reset accuracy metrics\n",
    "        acc_top1.reset()\n",
    "        acc_top5.reset()\n",
    "        btic = time.time()\n",
    "        train_loss = 0\n",
    "        num_batch = len(train_data)\n",
    "        \n",
    "        # Check and perform learning rate decay\n",
    "        if lr_decay_period and epoch and epoch % lr_decay_period == 0:\n",
    "            trainer.set_learning_rate(trainer.learning_rate*lr_decay)\n",
    "        elif lr_decay_period == 0 and epoch == lr_decay_epoch[lr_decay_count]:\n",
    "            trainer.set_learning_rate(trainer.learning_rate*lr_decay)\n",
    "            lr_decay_count += 1\n",
    "        # Loop over batches in an epoch\n",
    "        for i, batch in enumerate(train_data):\n",
    "            # Load train batch\n",
    "            data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)\n",
    "            label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0)\n",
    "            label_smooth = label\n",
    "            # Perform forward pass\n",
    "            with ag.record():\n",
    "                outputs = [net(X) for X in data]\n",
    "                loss = [L(yhat, y) for yhat, y in zip(outputs, label_smooth)]\n",
    "            # Perform backward pass\n",
    "            for l in loss:\n",
    "                l.backward()\n",
    "            # Perform updates\n",
    "            trainer.step(batch_size)\n",
    "            # Update accuracy metrics\n",
    "            acc_top1.update(label, outputs)\n",
    "            acc_top5.update(label, outputs)\n",
    "            # Update loss\n",
    "            train_loss += sum([l.sum().asscalar() for l in loss])\n",
    "            # Log training progress (after each `log_interval` batches)\n",
    "            if log_interval and not (i+1)%log_interval:\n",
    "                _, top1 = acc_top1.get()\n",
    "                _, top5 = acc_top5.get()\n",
    "                err_top1, err_top5 = (1-top1, 1-top5)\n",
    "                logging.info('Epoch[%d] Batch [%d]\\tSpeed: %f samples/sec\\ttop1-err=%f\\ttop5-err=%f'%(\n",
    "                             epoch, i, batch_size*log_interval/(time.time()-btic), err_top1, err_top5))\n",
    "                btic = time.time()\n",
    "                \n",
    "        # Retrieve training errors and loss\n",
    "        _, top1 = acc_top1.get()\n",
    "        _, top5 = acc_top5.get()\n",
    "        err_top1, err_top5 = (1-top1, 1-top5)\n",
    "        train_loss /= num_batch * batch_size\n",
    "\n",
    "        # Compute validation errors\n",
    "        err_top1_val, err_top5_val = test(ctx, val_data)\n",
    "        # Update training history\n",
    "        train_history.update([err_top1, err_top5, err_top1_val, err_top5_val])\n",
    "        # Update plot\n",
    "        train_history.plot(['training-top1-err', 'validation-top1-err','training-top5-err', 'validation-top5-err'],\n",
    "                           save_path='%s/%s_top_error.png'%(save_plot_dir, model_name))\n",
    "        logging.info('[Epoch %d] training: err-top1=%f err-top5=%f loss=%f'%(epoch, err_top1, err_top5, train_loss))\n",
    "        logging.info('[Epoch %d] time cost: %f'%(epoch, time.time()-tic))\n",
    "        logging.info('[Epoch %d] validation: err-top1=%f err-top5=%f'%(epoch, err_top1_val, err_top5_val))\n",
    "    \n",
    "        # Save a snapshot of the best model - use net.export to get MXNet symbols and params\n",
    "        if err_top1_val < best_val_score and epoch > 50:\n",
    "            best_val_score = err_top1_val\n",
    "            net.export('%s/%.4f-imagenet-%s-best'%(save_dir, best_val_score, model_name), epoch)\n",
    "        # Save a snapshot of the model after each 'save_frequency' epochs\n",
    "        if save_frequency and save_dir and (epoch + 1) % save_frequency == 0:\n",
    "            net.export('%s/%.4f-imagenet-%s'%(save_dir, best_val_score, model_name), epoch)\n",
    "    # Save a snapshot of the model at the end of training\n",
    "    if save_frequency and save_dir:\n",
    "        net.export('%s/%.4f-imagenet-%s'%(save_dir, best_val_score, model_name), epochs-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train model\n",
    "* Run the cell below to start training\n",
    "* Logs are displayed in the cell output\n",
    "* An example run of 1 epoch is shown here\n",
    "* Once training completes, the symbols and params files are saved in the root folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Epoch[0] Batch [49]\tSpeed: 252.603294 samples/sec\ttop1-err=0.998281\ttop5-err=0.991992\n",
      "INFO:root:Epoch[0] Batch [99]\tSpeed: 349.047239 samples/sec\ttop1-err=0.996875\ttop5-err=0.986914\n",
      "INFO:root:Epoch[0] Batch [149]\tSpeed: 561.580055 samples/sec\ttop1-err=0.995768\ttop5-err=0.982422\n",
      "INFO:root:Epoch[0] Batch [199]\tSpeed: 374.183433 samples/sec\ttop1-err=0.994756\ttop5-err=0.977900\n",
      "INFO:root:Epoch[0] Batch [249]\tSpeed: 696.750858 samples/sec\ttop1-err=0.993594\ttop5-err=0.974391\n",
      "INFO:root:Epoch[0] Batch [299]\tSpeed: 448.550708 samples/sec\ttop1-err=0.992741\ttop5-err=0.970579\n",
      "INFO:root:Epoch[0] Batch [349]\tSpeed: 657.723046 samples/sec\ttop1-err=0.991708\ttop5-err=0.966903\n",
      "INFO:root:Epoch[0] Batch [399]\tSpeed: 545.634811 samples/sec\ttop1-err=0.990640\ttop5-err=0.963311\n",
      "INFO:root:Epoch[0] Batch [449]\tSpeed: 671.446248 samples/sec\ttop1-err=0.989510\ttop5-err=0.960022\n",
      "INFO:root:Epoch[0] Batch [499]\tSpeed: 755.356674 samples/sec\ttop1-err=0.988352\ttop5-err=0.956621\n",
      "INFO:root:Epoch[0] Batch [549]\tSpeed: 691.063609 samples/sec\ttop1-err=0.987124\ttop5-err=0.952955\n",
      "INFO:root:Epoch[0] Batch [599]\tSpeed: 918.181990 samples/sec\ttop1-err=0.985859\ttop5-err=0.949352\n",
      "INFO:root:Epoch[0] Batch [649]\tSpeed: 860.124778 samples/sec\ttop1-err=0.984691\ttop5-err=0.945947\n",
      "INFO:root:Epoch[0] Batch [699]\tSpeed: 940.958643 samples/sec\ttop1-err=0.983491\ttop5-err=0.942600\n",
      "INFO:root:Epoch[0] Batch [749]\tSpeed: 1075.782670 samples/sec\ttop1-err=0.982203\ttop5-err=0.939073\n",
      "INFO:root:Epoch[0] Batch [799]\tSpeed: 1246.127198 samples/sec\ttop1-err=0.981028\ttop5-err=0.935557\n",
      "INFO:root:Epoch[0] Batch [849]\tSpeed: 1296.864854 samples/sec\ttop1-err=0.979768\ttop5-err=0.932013\n",
      "INFO:root:Epoch[0] Batch [899]\tSpeed: 1231.886412 samples/sec\ttop1-err=0.978496\ttop5-err=0.928592\n",
      "INFO:root:Epoch[0] Batch [949]\tSpeed: 1449.699107 samples/sec\ttop1-err=0.977200\ttop5-err=0.925043\n",
      "INFO:root:Epoch[0] Batch [999]\tSpeed: 1379.325339 samples/sec\ttop1-err=0.975988\ttop5-err=0.921682\n",
      "INFO:root:Epoch[0] Batch [1049]\tSpeed: 1733.103910 samples/sec\ttop1-err=0.974645\ttop5-err=0.918317\n",
      "INFO:root:Epoch[0] Batch [1099]\tSpeed: 1294.581225 samples/sec\ttop1-err=0.973402\ttop5-err=0.914707\n",
      "INFO:root:Epoch[0] Batch [1149]\tSpeed: 1772.343415 samples/sec\ttop1-err=0.972227\ttop5-err=0.911491\n",
      "INFO:root:Epoch[0] Batch [1199]\tSpeed: 1463.033384 samples/sec\ttop1-err=0.971081\ttop5-err=0.908232\n",
      "INFO:root:Epoch[0] Batch [1249]\tSpeed: 1432.549199 samples/sec\ttop1-err=0.969822\ttop5-err=0.905075\n",
      "INFO:root:Epoch[0] Batch [1299]\tSpeed: 1568.169224 samples/sec\ttop1-err=0.968579\ttop5-err=0.901854\n",
      "INFO:root:Epoch[0] Batch [1349]\tSpeed: 1544.073387 samples/sec\ttop1-err=0.967354\ttop5-err=0.898666\n",
      "INFO:root:Epoch[0] Batch [1399]\tSpeed: 1441.059195 samples/sec\ttop1-err=0.966059\ttop5-err=0.895533\n",
      "INFO:root:Epoch[0] Batch [1449]\tSpeed: 1669.211205 samples/sec\ttop1-err=0.964786\ttop5-err=0.892244\n",
      "INFO:root:Epoch[0] Batch [1499]\tSpeed: 1311.339637 samples/sec\ttop1-err=0.963413\ttop5-err=0.889146\n",
      "INFO:root:Epoch[0] Batch [1549]\tSpeed: 1591.729484 samples/sec\ttop1-err=0.962160\ttop5-err=0.886021\n",
      "INFO:root:Epoch[0] Batch [1599]\tSpeed: 1421.214632 samples/sec\ttop1-err=0.960835\ttop5-err=0.882871\n",
      "INFO:root:Epoch[0] Batch [1649]\tSpeed: 1593.170427 samples/sec\ttop1-err=0.959512\ttop5-err=0.879878\n",
      "INFO:root:Epoch[0] Batch [1699]\tSpeed: 1558.754237 samples/sec\ttop1-err=0.958289\ttop5-err=0.876978\n",
      "INFO:root:Epoch[0] Batch [1749]\tSpeed: 1412.428006 samples/sec\ttop1-err=0.957144\ttop5-err=0.874213\n",
      "INFO:root:Epoch[0] Batch [1799]\tSpeed: 1634.056272 samples/sec\ttop1-err=0.955924\ttop5-err=0.871291\n",
      "INFO:root:Epoch[0] Batch [1849]\tSpeed: 1496.589282 samples/sec\ttop1-err=0.954576\ttop5-err=0.868342\n",
      "INFO:root:Epoch[0] Batch [1899]\tSpeed: 1386.387604 samples/sec\ttop1-err=0.953237\ttop5-err=0.865400\n",
      "INFO:root:Epoch[0] Batch [1949]\tSpeed: 1760.235240 samples/sec\ttop1-err=0.951994\ttop5-err=0.862475\n",
      "INFO:root:Epoch[0] Batch [1999]\tSpeed: 1266.353942 samples/sec\ttop1-err=0.950653\ttop5-err=0.859561\n",
      "INFO:root:Epoch[0] Batch [2049]\tSpeed: 1986.385151 samples/sec\ttop1-err=0.949335\ttop5-err=0.856764\n",
      "INFO:root:Epoch[0] Batch [2099]\tSpeed: 1486.284735 samples/sec\ttop1-err=0.947954\ttop5-err=0.853871\n",
      "INFO:root:Epoch[0] Batch [2149]\tSpeed: 1500.575434 samples/sec\ttop1-err=0.946617\ttop5-err=0.850946\n",
      "INFO:root:Epoch[0] Batch [2199]\tSpeed: 1531.607468 samples/sec\ttop1-err=0.945342\ttop5-err=0.848274\n",
      "INFO:root:Epoch[0] Batch [2249]\tSpeed: 1391.015252 samples/sec\ttop1-err=0.944004\ttop5-err=0.845447\n",
      "INFO:root:Epoch[0] Batch [2299]\tSpeed: 1848.863785 samples/sec\ttop1-err=0.942735\ttop5-err=0.842776\n",
      "INFO:root:Epoch[0] Batch [2349]\tSpeed: 1418.983373 samples/sec\ttop1-err=0.941479\ttop5-err=0.840032\n",
      "INFO:root:Epoch[0] Batch [2399]\tSpeed: 1599.456423 samples/sec\ttop1-err=0.940103\ttop5-err=0.837321\n",
      "INFO:root:Epoch[0] Batch [2449]\tSpeed: 1565.836765 samples/sec\ttop1-err=0.938842\ttop5-err=0.834654\n",
      "INFO:root:Epoch[0] Batch [2499]\tSpeed: 1613.621373 samples/sec\ttop1-err=0.937505\ttop5-err=0.831930\n",
      "INFO:root:[Epoch 0] training: err-top1=0.937448 err-top5=0.831813 loss=5.506481\n",
      "INFO:root:[Epoch 0] time cost: 1320.680112\n",
      "INFO:root:[Epoch 0] validation: err-top1=0.860680 err-top5=0.672240\n"
     ]
    }
   ],
   "source": [
    "def main():\n",
    "    net.hybridize()\n",
    "    train(num_epochs, context)\n",
    "    #net.export(model_name)\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export model to ONNX format\n",
    "The conversion of the model to ONNX format is done using an internal converter which will be released soon. The notebook will be updated with the code for the export once the converter is released."
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "display_name": "",
  "kernelspec": {
   "display_name": "Environment (conda_anaconda3)",
   "language": "python",
   "name": "conda_anaconda3"
  },
  "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.4"
  },
  "name": ""
 },
 "nbformat": 4,
 "nbformat_minor": 2
}