vision/mobilenet_start.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "mobilenet_start.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python2",
      "display_name": "Python 2"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "9J7p406abzgl",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Part III: Separable convolutions  [Mobilenet](https://arxiv.org/pdf/1704.04861.pdf) \n",
        "- Modify the baseline model in part I to use *separable* convolutions, similar to Mobilenet.\n",
        "- Check number of parameters and compare with the baseline.\n",
        "- Train classifier.\n",
        "\n",
        "\n",
        "### Separable convolutions\n",
        "\n",
        "For example, a 2D conv can be written as a sequence of 2 1D conv, i.e. \\begin{equation} y[m,n]=x[m,n]*h[m,n] = h_1[m]*(h_2[n]*x[m,n])\\end{equation}\n",
        "\n",
        "assuming $x$ is a 2D input signal, $h$ is a 2D filter that can be separated into 2 1D filters $h_1$ and $h_2$, and $y$ is the output of convolving $x$ with $h$.  \n",
        "\n",
        "\n",
        "Similarly for 3D case, we apply the separability between feature channel and spatial dimensions (as shown in the figure below on the left), i.e. \\begin{equation} y[m,n,p]=x[m,n, p]*h[m,n,p] = h_1[p]*(h_2[m,n]*x[m,n,p])\\end{equation}\n",
        "\n",
        "![alt text](https://github.com/tmlss2018/PracticalSessions/blob/master/assets/separable.png?raw=true)\n",
        "\n",
        "### Architecture\n",
        "- baseline from Part I, adapted for separable convolutions; note that this is not exactly the model in Mobilenet paper. \n",
        "- first conv layer stays the same, kernel = (3,3), stride=1, num_channels = 64\n",
        "- depthwise conv layers:\n",
        "\n",
        "\n",
        ">* channel_multiplier = 1\n",
        ">* strides = [1, 2, 1, 1, 2, 1, 1, 2, 1, 1]\n",
        "\n",
        "- 1x1 conv layers:\n",
        "\n",
        ">* num_channels = [128, 128, 128, 256, 256, 256, 512, 512, 512, 512]\n",
        ">* strides = 1\n",
        "\n",
        "- the depthwise and 1x1 conv layers are added after the first conv layer, alternating between the two types \n",
        "- padding: SAME (snt.SAME) for all layers\n",
        "- num_output_classes = 10\n",
        "- BatchNorm and ReLU are added after each conv layer as illustrated in the figure above\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "id": "FhWI4Pix5GJw",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Imports"
      ]
    },
    {
      "metadata": {
        "id": "na0VvPXmYKp1",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import math\n",
        "import time\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "# Don't forget to select GPU runtime environment in Runtime -> Change runtime type\n",
        "device_name = tf.test.gpu_device_name()\n",
        "if device_name != '/device:GPU:0':\n",
        "  raise SystemError('GPU device not found')\n",
        "print('Found GPU at: {}'.format(device_name))\n",
        "\n",
        "# we will use Sonnet on top of TF \n",
        "!pip install -q dm-sonnet\n",
        "import sonnet as snt\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "# Plotting library.\n",
        "from matplotlib import pyplot as plt\n",
        "import pylab as pl\n",
        "from IPython import display"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "1xlKHOLbhvY7",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Reset graph\n",
        "tf.reset_default_graph()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "8g16XweXs2Uq",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Download dataset to be used for training and testing\n",
        "- Cifar-10 equivalent of MNIST for natural RGB images\n",
        "- 60000 32x32 colour images in 10 classes: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck\n",
        "- train: 50000; test: 10000"
      ]
    },
    {
      "metadata": {
        "id": "1g_EOx07s1XZ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "cifar10 = tf.keras.datasets.cifar10\n",
        "# (down)load dataset\n",
        "(train_images, train_labels), (test_images, test_labels) = cifar10.load_data()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "iZofMjOuUEOF",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "cell_type": "code",
      "source": [
        "#@title Prepare the data for training and testing\n",
        "# define dimension of the batches to sample from the datasets\n",
        "BATCH_SIZE_TRAIN = 64 #@param\n",
        "BATCH_SIZE_TEST = 100 #@param\n",
        "\n",
        "# create Dataset objects using the data previously downloaded\n",
        "dataset_train = tf.data.Dataset.from_tensor_slices((train_images, train_labels))\n",
        "# we shuffle the data and sample repeatedly batches for training\n",
        "batched_dataset_train = dataset_train.shuffle(100000).repeat().batch(BATCH_SIZE_TRAIN)\n",
        "# create iterator to retrieve batches\n",
        "iterator_train = batched_dataset_train.make_one_shot_iterator()\n",
        "# get a training batch of images and labels\n",
        "(batch_train_images, batch_train_labels) = iterator_train.get_next()\n",
        "\n",
        "# we do the same for test dataset\n",
        "dataset_test = tf.data.Dataset.from_tensor_slices((test_images, test_labels))\n",
        "batched_dataset_test = dataset_test.repeat().batch(BATCH_SIZE_TEST)\n",
        "iterator_test = batched_dataset_test.make_one_shot_iterator() \n",
        "(batch_test_images, batch_test_labels) = iterator_test.get_next()\n",
        "\n",
        "# Squeeze labels and convert from uint8 to int32 - required below by the loss op\n",
        "batch_test_labels = tf.cast(tf.squeeze(batch_test_labels), tf.int32)\n",
        "batch_train_labels = tf.cast(tf.squeeze(batch_train_labels), tf.int32)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "_PS2GjTxRZx9",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "cell_type": "code",
      "source": [
        "#@title Preprocessing of data\n",
        "# Data augmentation used for train preprocessing\n",
        "# - scale image to [-1 , 1]\n",
        "# - get a random crop\n",
        "# - apply horizontal flip randomly\n",
        "\n",
        "def train_image_preprocess(h, w, random_flip=True):\n",
        "  \"\"\"Image processing required for training the model.\"\"\"\n",
        "  \n",
        "  def random_flip_left_right(image, flip_index, seed=None):\n",
        "    shape = image.get_shape()\n",
        "    if shape.ndims == 3 or shape.ndims is None:\n",
        "      uniform_random = tf.random_uniform([], 0, 1.0, seed=seed)\n",
        "      mirror_cond = tf.less(uniform_random, .5)\n",
        "      result = tf.cond(\n",
        "          mirror_cond,\n",
        "          lambda: tf.reverse(image, [flip_index]),\n",
        "          lambda: image\n",
        "      )\n",
        "      return result\n",
        "    elif shape.ndims == 4:\n",
        "      uniform_random = tf.random_uniform(\n",
        "          [tf.shape(image)[0]], 0, 1.0, seed=seed\n",
        "      )\n",
        "      mirror_cond = tf.less(uniform_random, .5)\n",
        "      return tf.where(\n",
        "          mirror_cond,\n",
        "          image,\n",
        "          tf.map_fn(lambda x: tf.reverse(x, [flip_index]), image, dtype=image.dtype)\n",
        "      )\n",
        "    else:\n",
        "      raise ValueError(\"\\'image\\' must have either 3 or 4 dimensions.\")\n",
        "\n",
        "  def fn(image):\n",
        "    # Ensure the data is in range [-1, 1].\n",
        "    image = tf.image.convert_image_dtype(image, dtype=tf.float32)\n",
        "    image = image * 2.0 - 1.0\n",
        "    # Randomly choose a (h, w, 3) patch.\n",
        "    image = tf.random_crop(image, size=(BATCH_SIZE_TRAIN, h, w, 3))\n",
        "    # Randomly flip the image.\n",
        "    image = random_flip_left_right(image, 2)\n",
        "    return image\n",
        "\n",
        "  return fn\n",
        "\n",
        "# Test preprocessing: only scale to [-1,1].\n",
        "def test_image_preprocess():\n",
        "  def fn(image):\n",
        "    image = tf.image.convert_image_dtype(image, dtype=tf.float32)\n",
        "    image = image * 2.0 - 1.0\n",
        "    return image\n",
        "  return fn"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "HXGbNHHJ-y2z",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Define classifier using separable convolutions"
      ]
    },
    {
      "metadata": {
        "id": "zrRe4zFP2BIE",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class Mobilenet(snt.AbstractModule):\n",
        "  \n",
        "  def __init__(self, num_classes, name=\"mobilenet\"):\n",
        "    super(Mobilenet, self).__init__(name=name)\n",
        "    #############\n",
        "    #           #\n",
        "    # YOUR CODE #\n",
        "    #           #\n",
        "    #############\n",
        "   \n",
        "  def _build(self, inputs, is_training=None, test_local_stats=False):\n",
        "    #############\n",
        "    #           #\n",
        "    # YOUR CODE #\n",
        "    #           #\n",
        "    #############\n",
        "\n",
        "    return logits"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "OD8IR90m_0-r",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "cell_type": "code",
      "source": [
        "#@title Get number of parameters in a scope by iterating through the trainable variables\n",
        "def get_num_params(scope):\n",
        "  total_parameters = 0\n",
        "  for variable in tf.trainable_variables(scope):\n",
        "    # shape is an array of tf.Dimension\n",
        "    shape = variable.get_shape()\n",
        "    variable_parameters = 1\n",
        "    for dim in shape:\n",
        "      variable_parameters *= dim.value\n",
        "    total_parameters += variable_parameters\n",
        "  return total_parameters"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "inX9OlHW5xWR",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Instantiate the model and connect to data \n"
      ]
    },
    {
      "metadata": {
        "id": "TZzlpO0oJFZy",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "cell_type": "code",
      "source": [
        "# First define the preprocessing ops for the train/test data\n",
        "crop_height = 24 #@param\n",
        "cropt_width = 24 #@param\n",
        "preprocess_fn_train = train_image_preprocess(crop_height, cropt_width)\n",
        "preprocess_fn_test = test_image_preprocess()\n",
        "\n",
        "num_classes = 10 #@param"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ceTUdZ6FuQIv",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Instantiate the model\n",
        "with tf.variable_scope(\"mobilenet\"):\n",
        "  mobilenet_model = Mobilenet(num_classes=10)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "a86sl1lLmOVc",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Get predictions from the model; use the corresponding preprocess ops and is_training flag\n",
        "predictions_mobilenet = mobilenet_model(preprocess_fn_train(batch_train_images), is_training=True)\n",
        "print (predictions_mobilenet)\n",
        "\n",
        "test_predictions_mobilenet = mobilenet_model(preprocess_fn_test(batch_test_images), is_training=False)\n",
        "print (test_predictions_mobilenet)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "GskFb_d0tG8T",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Get number of parameters and compare with baseline"
      ]
    },
    {
      "metadata": {
        "id": "uLQeBvODtGJ9",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Can you obtain this number by hand?\n",
        "print (\"Total number of parameters of Mobilenet model\")\n",
        "print (get_num_params(\"mobilenet\"))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "AGGc24z7vpZi",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "cell_type": "code",
      "source": [
        "# @title Setup training (same as for baseline)\n",
        "def get_loss(logits=None, labels=None):\n",
        "  # We reduce over batch dimension, to ensure the loss is a scalar.   \n",
        "  return tf.reduce_mean(\n",
        "      tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
        "          labels=labels, logits=logits))\n",
        "\n",
        "# Define train and test loss functions\n",
        "train_loss = get_loss(labels=batch_train_labels, logits=predictions_mobilenet)\n",
        "test_loss = get_loss(labels=batch_test_labels, logits=test_predictions_mobilenet)\n",
        "\n",
        "# For evaluation, we look at top_k_accuracy since it's easier to interpret; normally k=1 or k=5\n",
        "def top_k_accuracy(k, labels, logits):\n",
        "  in_top_k = tf.nn.in_top_k(predictions=tf.squeeze(logits), targets=labels, k=k)\n",
        "  return tf.reduce_mean(tf.cast(in_top_k, tf.float32))\n",
        "\n",
        "def get_optimizer(step):\n",
        "  \"\"\"Get the optimizer used for training.\"\"\"\n",
        "  lr_init = 0.1 # initial value for the learning rate\n",
        "  lr_schedule = (40e3, 60e3, 80e3) # after how many iterations to reduce the learning rate\n",
        "  lr_schedule = tf.to_int64(lr_schedule)\n",
        "  lr_factor = 0.1 # reduce learning rate by this factor\n",
        "  \n",
        "  \n",
        "  num_epochs = tf.reduce_sum(tf.to_float(step >= lr_schedule))\n",
        "  lr = lr_init * lr_factor**num_epochs\n",
        "\n",
        "  return tf.train.MomentumOptimizer(learning_rate=lr, momentum=0.9)\n",
        "\n",
        "# Create a global step that is incremented during training; useful for e.g. learning rate annealing\n",
        "global_step = tf.train.get_or_create_global_step()\n",
        "\n",
        "# instantiate the optimizer\n",
        "optimizer = get_optimizer(global_step)\n",
        "\n",
        "# Get training ops\n",
        "training_mobilenet_op = optimizer.minimize(train_loss, global_step)\n",
        "\n",
        "# Retrieve the update ops, which contain the moving average ops\n",
        "update_ops = tf.group(*tf.get_collection(tf.GraphKeys.UPDATE_OPS))\n",
        "\n",
        "# Manually add the update ops to the dependency path executed at each training iteration\n",
        "training_mobilenet_op = tf.group(training_mobilenet_op, update_ops)\n",
        "\n",
        "# Get test ops\n",
        "test_acc_mobilenet_op = top_k_accuracy(1, batch_test_labels, test_predictions_mobilenet)\n",
        "\n",
        "# Function that takes a list of losses and plots them.\n",
        "def plot_losses(loss_list, steps):\n",
        "  display.clear_output(wait=True)\n",
        "  display.display(pl.gcf())\n",
        "  pl.plot(steps, loss_list, c='b')\n",
        "  time.sleep(1.0)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "mKOKrcm07DWW",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Training params\n"
      ]
    },
    {
      "metadata": {
        "id": "MWYjSf0rw4Ik",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "cell_type": "code",
      "source": [
        "# Define number of training iterations and reporting intervals\n",
        "TRAIN_ITERS = 90e3 #@param\n",
        "REPORT_TRAIN_EVERY = 10 #@param\n",
        "PLOT_EVERY = 500 #@param\n",
        "REPORT_TEST_EVERY = 1000 #@param\n",
        "TEST_ITERS = 50 #@param"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "m42U8EYn7YVB",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Train the model. Full training gives ~87% accuracy."
      ]
    },
    {
      "metadata": {
        "id": "_zXWOk-sxpy8",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Create the session and initialize variables\n",
        "sess = tf.Session()\n",
        "sess.run(tf.global_variables_initializer())\n",
        "\n",
        "train_iter = 0\n",
        "losses = []\n",
        "steps = []\n",
        "for train_iter in range(int(TRAIN_ITERS)):\n",
        "  _, train_loss_np = sess.run([training_mobilenet_op, train_loss])\n",
        "  \n",
        "  if (train_iter % REPORT_TRAIN_EVERY) == 0:\n",
        "    losses.append(train_loss_np)\n",
        "    steps.append(train_iter)\n",
        "  if (train_iter % PLOT_EVERY) == 0:\n",
        "    plot_losses(losses, steps)    \n",
        "    \n",
        "  if (train_iter % REPORT_TEST_EVERY) == 0:\n",
        "    avg_acc = 0.0\n",
        "    for test_iter in range(TEST_ITERS):\n",
        "      acc = sess.run(test_acc_mobilenet_op)\n",
        "      avg_acc += acc\n",
        "      \n",
        "    avg_acc /= (TEST_ITERS)\n",
        "    print ('Test acc at iter {0:5d} out of {1:5d} is {2:.2f}%'.format(int(train_iter), int(TRAIN_ITERS), avg_acc*100.0))"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}