06_Batchnorm_Folding.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantized Neural Network\n",
    "## 06 Batchnorm Folding\n",
    "by [Soon Yau Cheong](http://www.linkedin.com/in/soonyau)\n",
    "\n",
    "Alright, we now have a trained graph with quantized weights which we could export together with the quantization parameters i.e. offset and scale. However, there is one more optimization step just before we throw away the floating point weights, and that is batchnorm folding. Batchnorm folding is to merge the batchnorm operation into the convolutional layer's weights and hence removing the additional computation for batchnorm in inference time. We'll first go through batchnorm briefly, derive the equations and cross check with Tensorflow-Lite model\n",
    "\n",
    "### Batchnorm\n",
    "One of the difficulty in training deep neural network is that, after each weight updates, the statistics of the activations change. The next layer will now see input activations that have different statistics than it saw in last training step, and it can render the just updated weights to be useless, and this phenoma propagates through the network and make learning difficult. In 2015, Google researchers came out with an idea to tackle this problem, that is to normalize the activation to have zero mean and variance of 1, hence keeping the statistics almost same and allow more flexible weights initization. \n",
    "\n",
    "![alt text](images/batchnorm.png)\n",
    "\n",
    "Figure above shows the algorithm from the original paper \n",
    "[here](https://arxiv.org/abs/1502.03167). It is actually quite simple, first we calculate the mean and variance by taking moving average from mini batches. Then the activations are normalized, note that epsilon is small number to for numerical stability, to avoid dividing by zero if variance is close to zero. In Tensorflow, 0.001 is used as the epsilon. Then there are two learnable variables, $\\gamma$ and $\\beta$ for scale and offset. It is quite a lot of arithmetic operations here but we can remove all of them. How? Well, after the graph is frozen, all the variables (mean, variance, gamma and beta) became constant, and we'll take advantage of that and merge them into convolutional (or fully connected) layer that preceeds the batchnorm layer.\n",
    "\n",
    "### Folding\n",
    "Without loss of generality,  we'll use a simplified dot product equation and bias in place of convolutional operations to derive the equations, for every filter with weight w and bias b, the activation, y is\n",
    "\\begin{equation}\n",
    "y = (\\sum_Nw_ix_i)+b\n",
    "\\end{equation}\n",
    "Now we apply batchnorm to it\n",
    "\\begin{equation}\n",
    "y_{bn} = \\gamma \\hat{y} + \\beta \\\\\n",
    "y_{bn} = \\frac{\\gamma}{\\sqrt{\\sigma^2+\\epsilon}} (y-\\mu) + \\beta \\end{equation}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{equation}\n",
    "y_{bn} =  \\gamma'((\\sum_Nw_ix_i)+b-\\mu) + \\beta \\\\\n",
    "\\end{equation}\n",
    "\n",
    "where\n",
    "\\begin{equation}\n",
    "\\gamma' = \\frac{\\gamma}{\\sqrt{\\sigma^2+\\epsilon}}\n",
    "\\end{equation}\n",
    "Let's rearrange the variables\n",
    "\\begin{equation}\n",
    "y_{bn} =  \\sum_N (\\gamma' w_i)x_i+ \\gamma'(b-\\mu) + \\beta \\\\\n",
    "\\end{equation}\n",
    "\n",
    "Now, the new convolutional equations become\n",
    "\\begin{equation}\n",
    "y_{bn} =  \\sum_N \\hat{w_i}x_i+ \\hat{\\beta}\\\\\n",
    "\\end{equation}\n",
    "where\n",
    "\\begin{equation}\n",
    "\\hat{w_i} = (\\gamma' w_i)\n",
    "\\end{equation}\n",
    "\\begin{equation}\n",
    "\\hat{\\beta} = \\gamma'(b-\\mu) + \\beta \n",
    "\\end{equation}\n",
    "\n",
    "Now, the batchnorm parameters have all been folded into convolutional layer's weights and biases. As these are pre-calculated when exporting the weights, hence batchnorm layer is removed with no computational cost at all during inference time.\n",
    "\n",
    "### Implementation\n",
    "We can now load the weights from full precision graph, perform the batchnorm folding manually and see if we get the same result as the TFLite model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensorflow 1.10.0\n",
      "Python 3.5.2 (default, Nov 12 2018, 13:43:14) \n",
      "[GCC 5.4.0 20160609]\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.python import pywrap_tensorflow\n",
    "\n",
    "import utils\n",
    "\n",
    "print(\"Tensorflow\", tf.__version__)\n",
    "print(\"Python\", sys.version)\n",
    "\n",
    "# We now load the weight and bias from TFLite model\n",
    "model_path = 'models/mobilenet_v1/mobilenet_v1_1.0_224_quant.tflite'\n",
    "interpreter = tf.contrib.lite.Interpreter(model_path=model_path)                                         \n",
    "interpreter.allocate_tensors()\n",
    "\n",
    "# We load weights and biases of conv2d_0\n",
    "# where we get the index from Tutorial 1.\n",
    "tensor_idx = 8\n",
    "conv2d_0_w={}\n",
    "conv2d_0_w['detail'] = interpreter._get_tensor_details(tensor_idx)\n",
    "conv2d_0_w['tensor'] = interpreter.get_tensor(tensor_idx)\n",
    "\n",
    "tensor_idx = 6\n",
    "conv2d_0_b={}\n",
    "conv2d_0_b['detail'] = interpreter._get_tensor_details(tensor_idx)\n",
    "conv2d_0_b['tensor'] = interpreter.get_tensor(tensor_idx)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class TFGraphReader():\n",
    "    # load full-precision graph \n",
    "    def __init__(self, path=\"models/mobilenet_v1/mobilenet_v1_1.0_224_quant.ckpt\"):\n",
    "        self.reader = pywrap_tensorflow.NewCheckpointReader(tf_ckpt_path)\n",
    "        self.tensors = sorted(self.reader.get_variable_to_shape_map())\n",
    "        self.eps = 1e-3\n",
    "        \n",
    "    def batchnorm_fold(self, layer_name=\"Conv2d_0\"):\n",
    "        # read tensors\n",
    "        beta = self.reader.get_tensor(\"MobilenetV1/%s/BatchNorm/beta\"%layer_name)\n",
    "        gamma = self.reader.get_tensor(\"MobilenetV1/%s/BatchNorm/gamma\"%layer_name)\n",
    "        moving_mean = self.reader.get_tensor(\"MobilenetV1/%s/BatchNorm/moving_mean\"%layer_name)\n",
    "        moving_variance = self.reader.get_tensor(\"MobilenetV1/%s/BatchNorm/moving_variance\"%layer_name)\n",
    "        \n",
    "        weights = self.reader.get_tensor(\"MobilenetV1/%s/weights\"%layer_name)\n",
    "                \n",
    "        try:\n",
    "            biases = self.reader.get_tensor(\"MobilenetV1/%s/biases\"%layer_name)\n",
    "        except:\n",
    "            biases = 0\n",
    "            \n",
    "        # perform folding\n",
    "        gamma_ = gamma/np.sqrt(moving_variance+self.eps)        \n",
    "        weight_ = gamma_*weights\n",
    "        beta_ = gamma_*(biases-moving_mean)+beta\n",
    "        \n",
    "        # If we were to calculate the weight quantization offset and scale\n",
    "        '''\n",
    "        w_min = self.reader.get_tensor(\"MobilenetV1/MobilenetV1/%s/weights_quant/min\"%layer_name)\n",
    "        w_max = self.reader.get_tensor(\"MobilenetV1/MobilenetV1/%s/weights_quant/max\"%layer_name)       \n",
    "        quant_steps = 254 if narrow_range else 255\n",
    "        w_scale = (w_max - w_min)/quant_steps\n",
    "        w_offset = round(255 - w_max/w_scale)\n",
    "        '''\n",
    "        \n",
    "        return weight_, beta_\n",
    "\n",
    "\n",
    "# apply batchnorm to Conv2d_0 to return new weights and biases with batchnorm folded into them\n",
    "full_graph = TFGraphReader()\n",
    "w, b = full_graph.batchnorm_fold(\"Conv2d_0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we perform 8-bit quantization on the new weights and biases. For weights quantization, we can calculate the scale and offset like shown in the commented out code above. However, for efficient computation reasons which I'll describe in next tutorial, the scale and offset of biases are worked out from scale of input activation, weights and output activation. Therefore, in here, we'll just use the  params from TFLite model to perform quantization and to compare that with the quantized TFLite model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0.00017052092880476266, 0)\n"
     ]
    }
   ],
   "source": [
    "def quantize(x, scale, offset):\n",
    "    return (np.round(x/scale + offset)).astype(np.int32)\n",
    "\n",
    "quant = conv2d_0_w['detail']['quantization']\n",
    "bn_w = quantize(w, quant[0], quant[1])\n",
    "\n",
    "quant = conv2d_0_b['detail']['quantization']\n",
    "print(quant)\n",
    "bn_b = quantize(b, quant[0], quant[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's compare the them. Remember that the tensor dimension are different where the output feature is last dimension in TF but first dimension in TFLite."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[151 151 151]\n",
      "  [151 150 151]\n",
      "  [151 150 150]]\n",
      "\n",
      " [[151 151 151]\n",
      "  [151 151 151]\n",
      "  [151 150 150]]\n",
      "\n",
      " [[151 151 151]\n",
      "  [151 151 150]\n",
      "  [151 151 151]]]\n"
     ]
    }
   ],
   "source": [
    "# Folded TF weights\n",
    "print(bn_w[:,:,:,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[151 151 151]\n",
      "  [151 150 151]\n",
      "  [151 150 150]]\n",
      "\n",
      " [[151 151 151]\n",
      "  [151 151 151]\n",
      "  [151 150 150]]\n",
      "\n",
      " [[151 151 151]\n",
      "  [151 151 150]\n",
      "  [151 151 151]]]\n"
     ]
    }
   ],
   "source": [
    "# and from TFlite.\n",
    "print(conv2d_0_w['tensor'][0,:,:,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ -7254  13465  -1591  -2488   9901  13359   1947  16203  -2165  -7399\n",
      "  -5250  13549  17637   9441   3877  17663  -7985  13809  12408  11717\n",
      "  -3441   -104 -13034  17888 -12487  15548  26765  -2599  14359  10137\n",
      "  -2149  20334]\n",
      "[ -7254  13465  -1591  -2488   9901  13359   1947  16203  -2165  -7399\n",
      "  -5250  13549  17637   9441   3877  17663  -7985  13809  12408  11717\n",
      "  -3441   -104 -13034  17888 -12487  15548  26765  -2599  14359  10137\n",
      "  -2149  20334]\n"
     ]
    }
   ],
   "source": [
    "# do the same for biases\n",
    "print(bn_b)\n",
    "print(conv2d_0_b['tensor'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's Next?\n",
    "\n",
    "Today we learned about batchnorm folding to remove the batchnorm layers in inference. Next we'll look at how to further reduce additional computational arises from dequantization-quantization in fixed-point integer processor."
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}