vision/classification/squeezenet/train_squeezenet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Notebook for SqueezeNet 1.0 and 1.1 on ImageNet Dataset\n",
    "\n",
    "## Overview\n",
    "Use this notebook to train a SqueezeNet model from scratch. **Make sure to have the ImageNet dataset prepared** according to the guidelines in the dataset section in [SqueezeNet 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_squeezenet.py`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import dependencies\n",
    "Verify that all dependencies are installed using the cell below. Continue if no errors encountered, warnings can be ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "\n",
    "import 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.8xlarge ec2 instance on AWS. It has 4 Nvidia Tesla V100 GPUs (16GB each) and Intel(R) Xeon(R) CPU E5-2686 v4 @ 2.30GHz with 32 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": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify model - squeezenet1.0 or squeezenet1.1\n",
    "model_name = 'squeezenet1.0'\n",
    "\n",
    "# training and validation pictures to use\n",
    "data_dir = '/home/ubuntu/imagenet/img_dataset'\n",
    "\n",
    "# training batch size per device (CPU/GPU)\n",
    "batch_size = 128\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",
    "num_epochs = 100\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 = '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 = 100\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 = '.'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model definition in Gluon\n",
    "\n",
    "The class `SqueezeNet` contains model definitions of both SqueezeNet models and the required model is retrieved using the relevant constructor function: `Squeezenet1_0` or `Squeezenet1_1`. \n",
    "\n",
    "`_make_fire` and `_make_fire_conv` are helper functions which define the fire modules unique to SqueezeNet ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"SqueezeNet, implemented in Gluon.\"\"\"\n",
    "__all__ = ['SqueezeNet', 'squeezenet1_0', 'squeezenet1_1']\n",
    "\n",
    "# Helpers\n",
    "def _make_fire(squeeze_channels, expand1x1_channels, expand3x3_channels):\n",
    "    out = nn.HybridSequential(prefix='')\n",
    "    out.add(_make_fire_conv(squeeze_channels, 1))\n",
    "\n",
    "    paths = HybridConcurrent(axis=1, prefix='')\n",
    "    paths.add(_make_fire_conv(expand1x1_channels, 1))\n",
    "    paths.add(_make_fire_conv(expand3x3_channels, 3, 1))\n",
    "    out.add(paths)\n",
    "\n",
    "    return out\n",
    "\n",
    "def _make_fire_conv(channels, kernel_size, padding=0):\n",
    "    out = nn.HybridSequential(prefix='')\n",
    "    out.add(nn.Conv2D(channels, kernel_size, padding=padding))\n",
    "    out.add(nn.Activation('relu'))\n",
    "    return out\n",
    "\n",
    "# Net\n",
    "class SqueezeNet(HybridBlock):\n",
    "    r\"\"\"SqueezeNet model from the `\"SqueezeNet: AlexNet-level accuracy with 50x fewer parameters\n",
    "    and <0.5MB model size\" <https://arxiv.org/abs/1602.07360>`_ paper.\n",
    "    SqueezeNet 1.1 model from the `official SqueezeNet repo\n",
    "    <https://github.com/DeepScale/SqueezeNet/tree/master/SqueezeNet_v1.1>`_.\n",
    "    SqueezeNet 1.1 has 2.4x less computation and slightly fewer parameters\n",
    "    than SqueezeNet 1.0, without sacrificing accuracy.\n",
    "    Parameters\n",
    "    ----------\n",
    "    version : str\n",
    "        Version of squeezenet. Options are '1.0', '1.1'.\n",
    "    classes : int, default 1000\n",
    "        Number of classification classes.\n",
    "    \"\"\"\n",
    "    def __init__(self, version, classes=1000, **kwargs):\n",
    "        super(SqueezeNet, self).__init__(**kwargs)\n",
    "        assert version in ['1.0', '1.1'], (\"Unsupported SqueezeNet version {version}:\"\n",
    "                                           \"1.0 or 1.1 expected\".format(version=version))\n",
    "        with self.name_scope():\n",
    "            self.features = nn.HybridSequential(prefix='')\n",
    "            if version == '1.0':\n",
    "                self.features.add(nn.Conv2D(96, kernel_size=7, strides=2))\n",
    "                self.features.add(nn.Activation('relu'))\n",
    "                self.features.add(nn.MaxPool2D(pool_size=3, strides=2, ceil_mode=True))\n",
    "                self.features.add(_make_fire(16, 64, 64))\n",
    "                self.features.add(_make_fire(16, 64, 64))\n",
    "                self.features.add(_make_fire(32, 128, 128))\n",
    "                self.features.add(nn.MaxPool2D(pool_size=3, strides=2, ceil_mode=True))\n",
    "                self.features.add(_make_fire(32, 128, 128))\n",
    "                self.features.add(_make_fire(48, 192, 192))\n",
    "                self.features.add(_make_fire(48, 192, 192))\n",
    "                self.features.add(_make_fire(64, 256, 256))\n",
    "                self.features.add(nn.MaxPool2D(pool_size=3, strides=2, ceil_mode=True))\n",
    "                self.features.add(_make_fire(64, 256, 256))\n",
    "            else:\n",
    "                self.features.add(nn.Conv2D(64, kernel_size=3, strides=2))\n",
    "                self.features.add(nn.Activation('relu'))\n",
    "                self.features.add(nn.MaxPool2D(pool_size=3, strides=2, ceil_mode=True))\n",
    "                self.features.add(_make_fire(16, 64, 64))\n",
    "                self.features.add(_make_fire(16, 64, 64))\n",
    "                self.features.add(nn.MaxPool2D(pool_size=3, strides=2, ceil_mode=True))\n",
    "                self.features.add(_make_fire(32, 128, 128))\n",
    "                self.features.add(_make_fire(32, 128, 128))\n",
    "                self.features.add(nn.MaxPool2D(pool_size=3, strides=2, ceil_mode=True))\n",
    "                self.features.add(_make_fire(48, 192, 192))\n",
    "                self.features.add(_make_fire(48, 192, 192))\n",
    "                self.features.add(_make_fire(64, 256, 256))\n",
    "                self.features.add(_make_fire(64, 256, 256))\n",
    "            self.features.add(nn.Dropout(0.5))\n",
    "\n",
    "            self.output = nn.HybridSequential(prefix='')\n",
    "            self.output.add(nn.Conv2D(classes, kernel_size=1))\n",
    "            self.output.add(nn.Activation('relu'))\n",
    "            self.output.add(nn.AvgPool2D(13))\n",
    "            self.output.add(nn.Flatten())\n",
    "\n",
    "    def hybrid_forward(self, F, x):\n",
    "        x = self.features(x)\n",
    "        x = self.output(x)\n",
    "        return x\n",
    "\n",
    "# Constructor\n",
    "def get_squeezenet(version, root=os.path.join('~', '.mxnet', 'models'), **kwargs):\n",
    "    r\"\"\"SqueezeNet model from the `\"SqueezeNet: AlexNet-level accuracy with 50x fewer parameters\n",
    "    and <0.5MB model size\" <https://arxiv.org/abs/1602.07360>`_ paper.\n",
    "    SqueezeNet 1.1 model from the `official SqueezeNet repo\n",
    "    <https://github.com/DeepScale/SqueezeNet/tree/master/SqueezeNet_v1.1>`_.\n",
    "    SqueezeNet 1.1 has 2.4x less computation and slightly fewer parameters\n",
    "    than SqueezeNet 1.0, without sacrificing accuracy.\n",
    "    Parameters\n",
    "    ----------\n",
    "    version : str\n",
    "        Version of squeezenet. Options are '1.0', '1.1'.\n",
    "    root : str, default '~/.mxnet/models'\n",
    "        Location for keeping the model parameters.\n",
    "    \"\"\"\n",
    "    net = SqueezeNet(version, **kwargs)\n",
    "    return net\n",
    "\n",
    "def squeezenet1_0(**kwargs):\n",
    "    r\"\"\"SqueezeNet 1.0 model from the `\"SqueezeNet: AlexNet-level accuracy with 50x fewer parameters\n",
    "    and <0.5MB model size\" <https://arxiv.org/abs/1602.07360>`_ paper.\n",
    "    Parameters\n",
    "    ----------\n",
    "    root : str, default '~/.mxnet/models'\n",
    "        Location for keeping the model parameters.\n",
    "    \"\"\"\n",
    "    return get_squeezenet('1.0', **kwargs)\n",
    "\n",
    "def squeezenet1_1(**kwargs):\n",
    "    r\"\"\"SqueezeNet 1.1 model from the `official SqueezeNet repo\n",
    "    <https://github.com/DeepScale/SqueezeNet/tree/master/SqueezeNet_v1.1>`_.\n",
    "    SqueezeNet 1.1 has 2.4x less computation and slightly fewer parameters\n",
    "    than SqueezeNet 1.0, without sacrificing accuracy.\n",
    "    Parameters\n",
    "    ----------\n",
    "    root : str, default '~/.mxnet/models'\n",
    "        Location for keeping the model parameters.\n",
    "    \"\"\"\n",
    "    return get_squeezenet('1.1', **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper code\n",
    "Define context, optimizer, accuracy metrics, retireve gluon model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify logging fucntion\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",
    "\n",
    "# Retrieve gluon model\n",
    "if model_name == 'squeezenet1.0':\n",
    "    net = squeezenet1_0(**kwargs)\n",
    "else:\n",
    "    net = squeezenet1_1(**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",
    "\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": 13,
   "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.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": 14,
   "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": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train function\n",
    "def train(epochs, ctx):\n",
    "    if isinstance(ctx, mx.Context):\n",
    "        ctx = [ctx]\n",
    "    # Initialize network - Use method in MSRA paper <https://arxiv.org/abs/1502.01852>\n",
    "    net.initialize(mx.init.MSRAPrelu(), 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",
    "    # 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",
    "            ag.backward(loss)\n",
    "            # Perform updates\n",
    "            trainer.step(batch_size)\n",
    "            # Initialize last conv layer weights with normal distribution after first forward pass\n",
    "            if i==0 and epoch==0:\n",
    "                new_classifier_w = mx.nd.random_normal(shape=(1000, 512, 1, 1), scale=0.01)\n",
    "                final_conv_layer_params = net.output[0].params\n",
    "                final_conv_layer_params.get('weight').set_data(new_classifier_w)\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",
    "        # 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",
    "        \n",
    "        # Log training progress (after each epoch)\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": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Epoch[0] Batch [99]\tSpeed: 959.683155 samples/sec\ttop1-err=0.998848\ttop5-err=0.994902\n",
      "INFO:root:Epoch[0] Batch [199]\tSpeed: 1386.922154 samples/sec\ttop1-err=0.998984\ttop5-err=0.994941\n",
      "INFO:root:Epoch[0] Batch [299]\tSpeed: 1363.848197 samples/sec\ttop1-err=0.999023\ttop5-err=0.994902\n",
      "INFO:root:Epoch[0] Batch [399]\tSpeed: 1373.770129 samples/sec\ttop1-err=0.998984\ttop5-err=0.994771\n",
      "INFO:root:Epoch[0] Batch [499]\tSpeed: 1348.144496 samples/sec\ttop1-err=0.999000\ttop5-err=0.994789\n",
      "INFO:root:Epoch[0] Batch [599]\tSpeed: 1430.259184 samples/sec\ttop1-err=0.999014\ttop5-err=0.994863\n",
      "INFO:root:Epoch[0] Batch [699]\tSpeed: 1657.620811 samples/sec\ttop1-err=0.999021\ttop5-err=0.994819\n",
      "INFO:root:Epoch[0] Batch [799]\tSpeed: 1758.937257 samples/sec\ttop1-err=0.998960\ttop5-err=0.994653\n",
      "INFO:root:Epoch[0] Batch [899]\tSpeed: 1706.385409 samples/sec\ttop1-err=0.998891\ttop5-err=0.994390\n",
      "INFO:root:Epoch[0] Batch [999]\tSpeed: 1640.463420 samples/sec\ttop1-err=0.998805\ttop5-err=0.994051\n",
      "INFO:root:Epoch[0] Batch [1099]\tSpeed: 1767.966890 samples/sec\ttop1-err=0.998766\ttop5-err=0.993839\n",
      "INFO:root:Epoch[0] Batch [1199]\tSpeed: 1785.254212 samples/sec\ttop1-err=0.998698\ttop5-err=0.993579\n",
      "INFO:root:Epoch[0] Batch [1299]\tSpeed: 1661.894895 samples/sec\ttop1-err=0.998681\ttop5-err=0.993362\n",
      "INFO:root:Epoch[0] Batch [1399]\tSpeed: 1728.477546 samples/sec\ttop1-err=0.998630\ttop5-err=0.993152\n",
      "INFO:root:Epoch[0] Batch [1499]\tSpeed: 1734.059458 samples/sec\ttop1-err=0.998594\ttop5-err=0.992926\n",
      "INFO:root:Epoch[0] Batch [1599]\tSpeed: 1755.445442 samples/sec\ttop1-err=0.998556\ttop5-err=0.992681\n",
      "INFO:root:Epoch[0] Batch [1699]\tSpeed: 1760.107632 samples/sec\ttop1-err=0.998503\ttop5-err=0.992455\n",
      "INFO:root:Epoch[0] Batch [1799]\tSpeed: 1725.827186 samples/sec\ttop1-err=0.998437\ttop5-err=0.992232\n",
      "INFO:root:Epoch[0] Batch [1899]\tSpeed: 1707.523990 samples/sec\ttop1-err=0.998365\ttop5-err=0.991990\n",
      "INFO:root:Epoch[0] Batch [1999]\tSpeed: 1754.866448 samples/sec\ttop1-err=0.998299\ttop5-err=0.991673\n",
      "INFO:root:Epoch[0] Batch [2099]\tSpeed: 1725.601640 samples/sec\ttop1-err=0.998221\ttop5-err=0.991364\n",
      "INFO:root:Epoch[0] Batch [2199]\tSpeed: 1765.317844 samples/sec\ttop1-err=0.998125\ttop5-err=0.991025\n",
      "INFO:root:Epoch[0] Batch [2299]\tSpeed: 2288.659738 samples/sec\ttop1-err=0.998024\ttop5-err=0.990664\n",
      "INFO:root:Epoch[0] Batch [2399]\tSpeed: 2747.229720 samples/sec\ttop1-err=0.997944\ttop5-err=0.990334\n",
      "INFO:root:Epoch[0] Batch [2499]\tSpeed: 2847.152422 samples/sec\ttop1-err=0.997866\ttop5-err=0.989984\n",
      "INFO:root:[Epoch 0] training: err-top1=0.997866 err-top5=0.989975 loss=6.842851\n",
      "INFO:root:[Epoch 0] time cost: 821.784680\n",
      "INFO:root:[Epoch 0] validation: err-top1=0.994580 err-top5=0.979360\n"
     ]
    }
   ],
   "source": [
    "def main():\n",
    "    net.hybridize()\n",
    "    train(num_epochs, context)\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": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}