vision/distillation_start.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "distillation_start.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python2",
      "display_name": "Python 2"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "lzkm9_53o1H4",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Part IV: [Distilling the knowledge](https://arxiv.org/pdf/1503.02531.pdf) from a (larger) teacher model\n",
        "- Simple technique used in different problems to improve accuracy of models when a more complex model (that has access to privileged information) is available, e.g. [paper1](https://arxiv.org/pdf/1802.05668.pdf), [paper2](https://arxiv.org/pdf/1511.06342.pdf), [paper3](https://arxiv.org/pdf/1806.03863.pdf).\n",
        "- Import an already trained baseline model to use as teacher\n",
        "- Use the smaller Mobilenet model as student\n",
        "- Add KL distillation loss between teacher and student\n",
        "- Train Mobilenet student classifier with this joint loss\n",
        "\n",
        "#### Exercise:\n",
        "- Fill in the code for the student loss and train.\n"
      ]
    },
    {
      "metadata": {
        "id": "W5f4X7Zbpat0",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "###  The total loss for the student is:\n",
        "\\begin{equation}\n",
        "\\mathcal{L} = \\mathcal{L}_{\\text{classif}} + \\lambda \\mathcal{L}_{\\text{distill}}\n",
        "\\end{equation}\n",
        "\n",
        "For classification loss we use the regular cross-entropy and for the distillation loss, we use Kullback-Leibler (KL) divergence. $\\lambda$ is a normalisation factor explained below.\n",
        "\n",
        "\n",
        "**Reminder**:\n",
        "\n",
        "Given two distributions $t$ and $s$, we define their cross-entropy over a given set as:\n",
        "\n",
        "$$H(t,s) = H(t) + \\text{KL}(t,s),$$\n",
        "\n",
        "where $H(t)$ is the entropy of $t$, i.e. $H(t) = \\sum_{i=1}^{N}t(x_i) \\cdot \\log t(x_i)$\n",
        "\n",
        "and $\\text{KL}(t,s)$ is the KL divergence between $t$ and $s$, i.e. $\\text{KL}(t,s) = \\sum_{i=1}^{N}t(x_i) \\cdot \\log \\frac{t(x_i)}{s(x_i)} . $\n",
        "\n",
        "However, in most cases of interest to us, $t$ is a constant (either ground truth labels or teacher predictions also considered as constant), so the entropy term can be ignored since its gradient is 0. \n",
        "\n",
        "Hence we can use cross-entropy $H(t,s)$ for both losses: \n",
        "- the mismatch between ground truth and student predictions. \n",
        "- the mismatch between teacher and student distributions.\n",
        "\n",
        "In the context of distillation, it is useful to also remember that the outputs of the network are logits, which we interpret as probabilities when passed through softmax:\n",
        "\n",
        "$$p_i^{(T)} =\\frac{\\exp{(\\text{logits}_i / T) }}{\\sum_j \\exp{(\\text{logits}_j / T) }}. $$\n",
        "\n",
        "$T$ is the softmax temperature usually set to 1. Setting it to a higher value (e.g. 3.0 or 5.0) smooths the output probability distribution, an effect desired in distillation. More precisely, we will use\n",
        "\n",
        "\\begin{equation}\n",
        "\\mathcal{L}_{\\text{distill}} = H(\\text{p}_{\\text{teacher}}^{(T)}, \\text{p}_{\\text{student}}^{(T)}),\n",
        "\\end{equation}\n",
        "\n",
        "**The normalisation factor** \n",
        "\n",
        "$\\lambda$ is a normalisation factor that ensures the gradients of the two loss terms are comparable in scale. Note that the gradients of the distill loss term scale as $\\frac{1}{T^2}$ due to the logits being divided by $T$.  Hence we use $$\\lambda = T^2$$ to bring distillation term gradients to the same scale as the classification term gradients.\n",
        "\n",
        "\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": "V5JKC1HMpnmF",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Copy the pretrained weights of baseline model on the virtual machine\n",
        "- you need to load all three files in the *baseline* folder"
      ]
    },
    {
      "metadata": {
        "id": "cubpPmHgECbc",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from google.colab import files\n",
        "\n",
        "uploaded = files.upload()\n",
        "\n",
        "for fn in uploaded.keys():\n",
        "  print('User uploaded file \"{name}\" with length {length} bytes'.format(\n",
        "      name=fn, length=len(uploaded[fn])))"
      ],
      "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()\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "-7HJGeno09Rs",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print (train_labels.shape)\n",
        "print (test_labels.shape)"
      ],
      "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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "jjLIoHEI1L9u",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print (batch_test_labels)\n",
        "print (batch_train_labels)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "S-FnScIy1ITP",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# 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": "2scBoc09ZsO4",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "cell_type": "code",
      "source": [
        "#@title Teacher model is the baseline\n",
        "class Baseline(snt.AbstractModule):\n",
        "  \n",
        "  def __init__(self, num_classes, name=\"baseline\"):\n",
        "    super(Baseline, self).__init__(name=name)\n",
        "    self._num_classes = num_classes\n",
        "    self._output_channels = [\n",
        "        64, 64, 128, 128, 128, 256, 256, 256, 512, 512, 512\n",
        "        ]\n",
        "    self._num_layers = len(self._output_channels)\n",
        "\n",
        "    self._kernel_shapes = [[3, 3]] * self._num_layers  # All kernels are 3x3.\n",
        "    self._strides = [1, 1, 2, 1, 1, 2, 1, 1, 2, 1, 1]\n",
        "    self._paddings = [snt.SAME] * self._num_layers\n",
        "   \n",
        "  def _build(self, inputs, is_training=None, test_local_stats=False):\n",
        "    net = inputs\n",
        "    # instantiate all the convolutional layers \n",
        "    layers = [snt.Conv2D(name=\"conv_2d_{}\".format(i),\n",
        "                         output_channels=self._output_channels[i],\n",
        "                         kernel_shape=self._kernel_shapes[i],\n",
        "                         stride=self._strides[i],\n",
        "                         padding=self._paddings[i],\n",
        "                         use_bias=True) for i in xrange(self._num_layers)]\n",
        "    # connect them to the graph, adding batch norm and non-linearity\n",
        "    for i, layer in enumerate(layers):\n",
        "      net = layer(net)\n",
        "      bn = snt.BatchNorm(name=\"batch_norm_{}\".format(i))\n",
        "      net = bn(net, is_training=is_training, test_local_stats=test_local_stats)\n",
        "      net = tf.nn.relu(net)\n",
        "\n",
        "    net = tf.reduce_mean(net, reduction_indices=[1, 2], keepdims=False,\n",
        "                         name=\"avg_pool\")\n",
        "\n",
        "    logits = snt.Linear(self._num_classes)(net)\n",
        "\n",
        "    return logits"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "vR3pnr5NVDwj",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "cell_type": "code",
      "source": [
        "#@title Student model is Mobilenet\n",
        "class Mobilenet(snt.AbstractModule):\n",
        "  \n",
        "  def __init__(self, num_classes, name=\"mobilenet\"):\n",
        "    super(Mobilenet, self).__init__(name=name)\n",
        "    self._num_classes = num_classes\n",
        "    self._output_channels_first_conv = 64\n",
        "    self._output_channels_1x1 = [\n",
        "        128, 128, 128, 256, 256, 256, 512, 512, 512, 512\n",
        "    ]\n",
        "    self._strides_dw = [1, 2, 1, 1, 2, 1, 1, 2, 1, 1]\n",
        "    self._num_layers_dw = len(self._strides_dw)\n",
        "    self._num_layers_1x1 = len(self._output_channels_1x1)\n",
        "   \n",
        "  def _build(self, inputs, is_training=None, test_local_stats=False):\n",
        "    net = inputs\n",
        "    # instantiate all the convolutional layers\n",
        "    first_conv = snt.Conv2D(name=\"conv_2d_0\",\n",
        "                            output_channels=self._output_channels_first_conv,\n",
        "                            kernel_shape=3,\n",
        "                            stride=1,\n",
        "                            padding=snt.SAME,\n",
        "                            use_bias=True)\n",
        "    \n",
        "    # instantiate depthwise conv layers\n",
        "    conv_layers_dw = [snt.DepthwiseConv2D(name=\"conv_dw_2d_{}\".format(i),\n",
        "                                          channel_multiplier=1,\n",
        "                                          kernel_shape=3,\n",
        "                                          stride=self._strides_dw[i],\n",
        "                                          padding=snt.SAME,\n",
        "                                          use_bias=True)\n",
        "                      for i in xrange(self._num_layers_dw)]\n",
        "    \n",
        "    # instantiate 1x1 conv layers\n",
        "    conv_layers_1x1 = [snt.Conv2D(name=\"conv_1x1_2d_{}\".format(i),\n",
        "                                  output_channels=self._output_channels_1x1[i],\n",
        "                                  kernel_shape=1,\n",
        "                                  stride=1,\n",
        "                                  padding=snt.SAME,\n",
        "                                  use_bias=True)\n",
        "                       for i in xrange(self._num_layers_1x1)]\n",
        "    # connect first layer to the graph, adding batch norm and non-linearity\n",
        "    net = first_conv(net)\n",
        "    bn = snt.BatchNorm(name=\"batch_norm_0\")\n",
        "    net = bn(net, is_training=is_training, test_local_stats=test_local_stats)\n",
        "    net = tf.nn.relu(net)\n",
        "    \n",
        "    # connect the rest of the layers\n",
        "    for i, (layer_dw, layer_1x1) in enumerate(zip(conv_layers_dw, conv_layers_1x1)):\n",
        "      net = layer_dw(net)\n",
        "      bn = snt.BatchNorm(name=\"batch_norm_{}_0\".format(i))\n",
        "      net = bn(net, is_training=is_training, test_local_stats=test_local_stats)\n",
        "      net = tf.nn.relu(net)\n",
        "      net = layer_1x1(net)\n",
        "      bn = snt.BatchNorm(name=\"batch_norm_{}_1\".format(i))\n",
        "      net = bn(net, is_training=is_training, test_local_stats=test_local_stats)\n",
        "      net = tf.nn.relu(net)      \n",
        "\n",
        "    net = tf.reduce_mean(net, reduction_indices=[1, 2], keepdims=False,\n",
        "                         name=\"avg_pool\")\n",
        "\n",
        "    logits = snt.Linear(self._num_classes)(net)\n",
        "\n",
        "    return logits"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "mpFiYJDqSBeq",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Training params"
      ]
    },
    {
      "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",
        "# apply preprocess\n",
        "batch_train_images = preprocess_fn_train(batch_train_images)\n",
        "batch_test_images = preprocess_fn_test(batch_test_images)\n",
        "\n",
        "num_classes = 10 #@param\n",
        "\n",
        "# Define number of training iterations and reporting intervals\n",
        "TRAIN_ITERS = 120e3 #@param\n",
        "REPORT_TRAIN_EVERY = 100 #@param\n",
        "PLOT_EVERY = 500 #@param\n",
        "REPORT_TEST_EVERY = 1000 #@param\n",
        "TEST_ITERS = 50 #@param"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "698eQkBaVtNg",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# 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=tf.squeeze(tf.cast(labels, tf.int32)), k=k)\n",
        "  return tf.reduce_mean(tf.cast(in_top_k, tf.float32))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "0xI0ftMtqQjG",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Instantiate teacher and load pre-trained weights\n"
      ]
    },
    {
      "metadata": {
        "id": "V9eSqIp5WAKF",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "with tf.variable_scope(\"teacher\"):\n",
        "  teacher_model = Baseline(num_classes)\n",
        "predictions_teacher = teacher_model(batch_train_images, is_training=False)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "CtfBZ1_OO3L6",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### We do not want to alter the teacher weights, so apply `tf.stop_gradients` to `predictions_teacher`"
      ]
    },
    {
      "metadata": {
        "id": "S0hD3ZL2O04r",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "predictions_teacher = ############# YOUR CODE HERE #############"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "jtLgURi-O1gA",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Create saver to restore the pre-trained model\n",
        "# First remove the scope name from variables name, since the name in the checkpoint doesn't include it\n",
        "var_list = snt.get_variables_in_scope(\"teacher\", collection=tf.GraphKeys.GLOBAL_VARIABLES)  \n",
        "var_map = {}\n",
        "for i in range(0, len(var_list)):\n",
        "  name = var_list[i].name[len(\"teacher/\"):-2]\n",
        "  var_map[name] = var_list[i]\n",
        "  \n",
        "saver = tf.train.Saver(var_map, reshape=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "UhXzKr-TqlPx",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Instantiate student"
      ]
    },
    {
      "metadata": {
        "id": "g3S76cQUqksA",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "with tf.variable_scope(\"student\"):\n",
        "  student_model = Mobilenet(num_classes=num_classes)\n",
        "# get predictions from the model\n",
        "predictions_student = student_model(batch_train_images, is_training=True)\n",
        "test_predictions_student = student_model(batch_test_images, is_training=False)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "n18Yw0PtvX9i",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "cell_type": "code",
      "source": [
        "#@title Compare number of parameters between teacher and student\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\n",
        "\n",
        "print ('Number of paramters of teacher')\n",
        "print (get_num_params('teacher'))\n",
        "print ('Number of parameters of student')\n",
        "print (get_num_params('student'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "TjgjChm_rAm2",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Set up the training for student, adding the distillation loss weighted by the square of temperature as explained above. \n",
        "\n",
        "Normally we use T = 1, but for distillation we use T>1, e.g. T=5. We will visualise later the impact of T on logits."
      ]
    },
    {
      "metadata": {
        "id": "UcaKA4XfaE_q",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "cell_type": "code",
      "source": [
        "T_distill = 5.0 #@param \n",
        "T_normal = 1.0 #@param"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "2kQJkAIpZs5s",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#### First define the regular cross-entropy classification loss"
      ]
    },
    {
      "metadata": {
        "id": "vvnDmLsUc6JU",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def get_cross_entropy_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.softmax_cross_entropy_with_logits_v2(\n",
        "          labels=tf.one_hot(tf.squeeze(labels), num_classes), logits=logits))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "FRhpBqgOQPXa",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#### Define the distillation loss\n",
        "\n",
        "You can do this either with\n",
        "\n",
        "* `tf.distributions.kl_divergence` between the teacher and student distributions, respectively; or \n",
        "* `softmax_cross_entropy_with_logits`. Remember that in this case the labels are expected to sum to 1, while the output of the teacher network is logits. So we need to apply `softmax` on the `predictions_teacher`.\n"
      ]
    },
    {
      "metadata": {
        "id": "ivYOPAmkQOBy",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Using tf.distributions.kl_divergence\n",
        "#########################\n",
        "#                       #\n",
        "# YOUR CODE             #\n",
        "# distill_kl_loss = ... #         \n",
        "#########################\n",
        "\n",
        "# OR simpler, using tf.nn.softmax_cross_entropy_with_logits\n",
        "#########################\n",
        "#                       #\n",
        "# YOUR CODE             #\n",
        "# distill_kl_loss = ... #         \n",
        "#########################"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "nsfBoN1iQ8D3",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "**Define the joint training loss**"
      ]
    },
    {
      "metadata": {
        "id": "amN8bJDcQ6UQ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "lambda_ = ############## YOUR CODE ##############\n",
        "\n",
        "train_loss = get_cross_entropy_loss(logits=predictions_student, labels=batch_train_labels)\n",
        "train_loss += lambda_ * distill_kl_loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "bg5vcWv1S6Ne",
        "colab_type": "code",
        "colab": {},
        "cellView": "both"
      },
      "cell_type": "code",
      "source": [
        "#@title Set up the training; better to start with lower lr and longer training schedule\n",
        "def get_optimizer(step):\n",
        "  \"\"\"Get the optimizer used for training.\"\"\"\n",
        "  lr_schedule = (80e3, 100e3, 110e3)\n",
        "  lr_schedule = tf.to_int64(lr_schedule)\n",
        "  lr_factor = 0.1\n",
        "  \n",
        "  lr_init = 0.01\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, including BatchNorm update ops\n",
        "training_op = optimizer.minimize(train_loss, global_step)\n",
        "update_ops = tf.group(*tf.get_collection(tf.GraphKeys.UPDATE_OPS))\n",
        "training_op = tf.group(training_op, update_ops)\n",
        "\n",
        "# Display loss function\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": "BLzmNbeLSdgf",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Teacher and student accuracy"
      ]
    },
    {
      "metadata": {
        "id": "JbHwP67QSaGh",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "test_acc = top_k_accuracy(1, batch_test_labels, test_predictions_student)\n",
        "\n",
        "# We compute the accuracy of the teacher on the train set to make sure that\n",
        "# the loading of the pre-trained weights was successful\n",
        "acc_teacher = top_k_accuracy(1, batch_train_labels, predictions_teacher) "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ekWY7feta-mn",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Define ops to visualise the impact of softmax temperature on output distributions"
      ]
    },
    {
      "metadata": {
        "id": "pHh6LY61ZXhz",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "probs_high_temp = tf.nn.softmax(tf.div(predictions_teacher, T_distill)) \n",
        "probs_low_temp = tf.nn.softmax(tf.div(predictions_teacher, T_normal))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "tlnD9xw9bQG6",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Create the session and initialise variables"
      ]
    },
    {
      "metadata": {
        "id": "M6ajGgfzcdh3",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "sess = tf.Session()\n",
        "sess.run(tf.global_variables_initializer())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Q8DKMYb3rIg3",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Load pre-trained weights for teacher, and check accuracy to make sure the import was successful"
      ]
    },
    {
      "metadata": {
        "id": "LUUJ84-QbyOA",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "saver.restore(sess, \"baseline.ckpt\")\n",
        "\n",
        "num_batches = 100  # 100 batches * 64 samples per batch = 6400 out of 50000\n",
        "avg_accuracy = 0.0\n",
        "for _ in range(num_batches):\n",
        "  accuracy = sess.run(acc_teacher)\n",
        "  avg_accuracy += accuracy\n",
        "avg_accuracy /= num_batches\n",
        "\n",
        "# expected_accuracy > 0.90\n",
        "print (\"Teacher accuracy on a subset of the train set {:.3f}\".format(avg_accuracy))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "sIK7-OSdSkVY",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Visualize the impact of temperature on the logits"
      ]
    },
    {
      "metadata": {
        "id": "6GHU1c6bs6ir",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "probs_ht, probs_lt, gt = sess.run([probs_high_temp, probs_low_temp, tf.one_hot(batch_train_labels, num_classes)])\n",
        "# pick one sample and plot\n",
        "idx = 10\n",
        "plt.plot(probs_ht[idx], c='r', label='High Temp')\n",
        "plt.plot(probs_lt[idx], c='g', label='Low Temp')\n",
        "plt.plot(gt[idx], 'b--', label='GT')\n",
        "plt.xlim([0,9])\n",
        "plt.legend()\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "QHE75BhpSoVk",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Train the model. Full training gives ~92% accuracy.\n",
        "\n",
        "If running out of memory, reduce the BATCH_SIZE_TRAIN, e.g. 32 or 16.\n",
        "\n",
        "Note that the execution is slower and more memory is needed now, since for each training iteration of the student we need to run the forward pass for the teacher as well."
      ]
    },
    {
      "metadata": {
        "id": "C2877dV3Sxxr",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "train_iter = 0\n",
        "losses = []\n",
        "steps = []\n",
        "for train_iter in range(int(TRAIN_ITERS)):\n",
        "  _, train_loss_np = sess.run([training_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)\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": []
    }
  ]
}