vision/classification/vgg/train_vgg.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Notebook for VGG models on ImageNet Dataset\n",
    "\n",
    "## Overview\n",
    "Use this notebook to train a VGG model from scratch. **Make sure to have the ImageNet dataset prepared** according to the guidelines in the dataset section in [VGG 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_vgg.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.initializer import Xavier\n",
    "from mxnet.gluon.contrib.nn import HybridConcurrent\n",
    "from __future__ import division\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 vgg16, vgg16_bn, vgg19, vgg19_bn (_bn versions are with batch normalization)\n",
    "model_name = 'vgg16'\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 = 32\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.0001\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 = '50,80'\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 = 1000\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",
    "The class `VGG` contains model definitions of all VGG models and the required model is retrieved using the relevant constructor function: `vgg16` / `vgg16_bn` / `vgg19` / `vgg19_bn`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"VGG, implemented in Gluon.\"\"\"\n",
    "__all__ = ['VGG', 'vgg16', 'vgg19', 'vgg16_bn', 'vgg19_bn', 'get_vgg']\n",
    "\n",
    "\n",
    "class VGG(HybridBlock):\n",
    "    r\"\"\"VGG model from the `\"Very Deep Convolutional Networks for Large-Scale Image Recognition\"\n",
    "    <https://arxiv.org/abs/1409.1556>`_ paper.\n",
    "    Parameters\n",
    "    ----------\n",
    "    layers : list of int\n",
    "        Numbers of layers in each feature block.\n",
    "    filters : list of int\n",
    "        Numbers of filters in each feature block. List length should match the layers.\n",
    "    classes : int, default 1000\n",
    "        Number of classification classes.\n",
    "    batch_norm : bool, default False\n",
    "        Use batch normalization.\n",
    "    \"\"\"\n",
    "    def __init__(self, layers, filters, classes=1000, batch_norm=False, **kwargs):\n",
    "        super(VGG, self).__init__(**kwargs)\n",
    "        assert len(layers) == len(filters)\n",
    "        with self.name_scope():\n",
    "            self.features = self._make_features(layers, filters, batch_norm)\n",
    "            self.features.add(nn.Dense(4096, activation='relu',\n",
    "                                       weight_initializer='normal',\n",
    "                                       bias_initializer='zeros'))\n",
    "            self.features.add(nn.Dropout(rate=0.5))\n",
    "            self.features.add(nn.Dense(4096, activation='relu',\n",
    "                                       weight_initializer='normal',\n",
    "                                       bias_initializer='zeros'))\n",
    "            self.features.add(nn.Dropout(rate=0.5))\n",
    "            self.output = nn.Dense(classes,\n",
    "                                   weight_initializer='normal',\n",
    "                                   bias_initializer='zeros')\n",
    "\n",
    "    def _make_features(self, layers, filters, batch_norm):\n",
    "        featurizer = nn.HybridSequential(prefix='')\n",
    "        for i, num in enumerate(layers):\n",
    "            for _ in range(num):\n",
    "                featurizer.add(nn.Conv2D(filters[i], kernel_size=3, padding=1,\n",
    "                                         weight_initializer=Xavier(rnd_type='gaussian',\n",
    "                                                                   factor_type='out',\n",
    "                                                                   magnitude=2),\n",
    "                                         bias_initializer='zeros'))\n",
    "                if batch_norm:\n",
    "                    featurizer.add(nn.BatchNorm())\n",
    "                featurizer.add(nn.Activation('relu'))\n",
    "            featurizer.add(nn.MaxPool2D(strides=2))\n",
    "        return featurizer\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",
    "vgg_spec = {16: ([2, 2, 3, 3, 3], [64, 128, 256, 512, 512]),\n",
    "            19: ([2, 2, 4, 4, 4], [64, 128, 256, 512, 512])}\n",
    "\n",
    "\n",
    "# Constructors\n",
    "def get_vgg(num_layers, root=os.path.join('~', '.mxnet', 'models'), **kwargs):\n",
    "    r\"\"\"VGG model from the `\"Very Deep Convolutional Networks for Large-Scale Image Recognition\"\n",
    "    <https://arxiv.org/abs/1409.1556>`_ paper.\n",
    "    Parameters\n",
    "    ----------\n",
    "    num_layers : int\n",
    "        Number of layers for the variant of densenet. Options are 16, 19.\n",
    "    root : str, default '~/.mxnet/models'\n",
    "        Location for keeping the model parameters.\n",
    "    \"\"\"\n",
    "    layers, filters = vgg_spec[num_layers]\n",
    "    net = VGG(layers, filters, **kwargs)\n",
    "    return net\n",
    "\n",
    "def vgg16(**kwargs):\n",
    "    r\"\"\"VGG-16 model from the `\"Very Deep Convolutional Networks for Large-Scale Image Recognition\"\n",
    "    <https://arxiv.org/abs/1409.1556>`_ paper.\n",
    "    Parameters\n",
    "    ----------\n",
    "    root : str, default '~/.mxnet/models'\n",
    "        Location for keeping the model parameters.\n",
    "    \"\"\"\n",
    "    return get_vgg(16, **kwargs)\n",
    "\n",
    "def vgg19(**kwargs):\n",
    "    r\"\"\"VGG-19 model from the `\"Very Deep Convolutional Networks for Large-Scale Image Recognition\"\n",
    "    <https://arxiv.org/abs/1409.1556>`_ paper.\n",
    "    Parameters\n",
    "    ----------\n",
    "    root : str, default '~/.mxnet/models'\n",
    "        Location for keeping the model parameters.\n",
    "    \"\"\"\n",
    "    return get_vgg(19, **kwargs)\n",
    "\n",
    "def vgg16_bn(**kwargs):\n",
    "    r\"\"\"VGG-16 model with batch normalization from the\n",
    "    `\"Very Deep Convolutional Networks for Large-Scale Image Recognition\"\n",
    "    <https://arxiv.org/abs/1409.1556>`_ paper.\n",
    "    Parameters\n",
    "    ----------\n",
    "    root : str, default '~/.mxnet/models'\n",
    "        Location for keeping the model parameters.\n",
    "    \"\"\"\n",
    "    kwargs['batch_norm'] = True\n",
    "    return get_vgg(16, **kwargs)\n",
    "\n",
    "def vgg19_bn(**kwargs):\n",
    "    r\"\"\"VGG-19 model with batch normalization from the\n",
    "    `\"Very Deep Convolutional Networks for Large-Scale Image Recognition\"\n",
    "    <https://arxiv.org/abs/1409.1556>`_ paper.\n",
    "    Parameters\n",
    "    ----------\n",
    "    root : str, default '~/.mxnet/models'\n",
    "        Location for keeping the model parameters.\n",
    "    \"\"\"\n",
    "    kwargs['batch_norm'] = True\n",
    "    return get_vgg(19, **kwargs)\n",
    "\n",
    "models = {    'vgg16': vgg16,\n",
    "              'vgg19': vgg19,\n",
    "              'vgg16_bn': vgg16_bn,\n",
    "              'vgg19_bn': vgg19_bn\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, 'batch_norm': False}\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",
    "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",
    "\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.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": [
    "# 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",
    "            # 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",
    "\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": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Epoch[0] Batch [999]\tSpeed: 631.525741 samples/sec\ttop1-err=0.998711\ttop5-err=0.993820\n",
      "INFO:root:Epoch[0] Batch [1999]\tSpeed: 693.004984 samples/sec\ttop1-err=0.998566\ttop5-err=0.992742\n",
      "INFO:root:Epoch[0] Batch [2999]\tSpeed: 692.148446 samples/sec\ttop1-err=0.998234\ttop5-err=0.991396\n",
      "INFO:root:Epoch[0] Batch [3999]\tSpeed: 692.413054 samples/sec\ttop1-err=0.997648\ttop5-err=0.988691\n",
      "INFO:root:Epoch[0] Batch [4999]\tSpeed: 693.148666 samples/sec\ttop1-err=0.996670\ttop5-err=0.984130\n",
      "INFO:root:Epoch[0] Batch [5999]\tSpeed: 691.513520 samples/sec\ttop1-err=0.995020\ttop5-err=0.977861\n",
      "INFO:root:Epoch[0] Batch [6999]\tSpeed: 692.626715 samples/sec\ttop1-err=0.992588\ttop5-err=0.969491\n",
      "INFO:root:Epoch[0] Batch [7999]\tSpeed: 692.031051 samples/sec\ttop1-err=0.989684\ttop5-err=0.959757\n",
      "INFO:root:Epoch[0] Batch [8999]\tSpeed: 692.713621 samples/sec\ttop1-err=0.986289\ttop5-err=0.949328\n",
      "INFO:root:Epoch[0] Batch [9999]\tSpeed: 692.846073 samples/sec\ttop1-err=0.982526\ttop5-err=0.938502\n",
      "INFO:root:[Epoch 0] training: err-top1=0.982482 err-top5=0.938400 loss=6.210697\n",
      "INFO:root:[Epoch 0] time cost: 1913.175402\n",
      "INFO:root:[Epoch 0] validation: err-top1=0.921820 err-top5=0.776760\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": "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}