03_Quantizing_Tensorflow_Graph.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantized Neural Network \n",
    "## 03- Quantizing Tensorflow Graph\n",
    "\n",
    "by [Soon Yau Cheong](http://www.linkedin.com/in/soonyau)\n",
    "\n",
    "In last tutorial, we learned about quantization and dequantization of tensor, and also introduced TF function tf.fake_quant_with_min_max_args to do that. Today, we will first look at how to quantize a pre-trained Tensorflow graph, followed by creating a quantization-aware graph for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import utils\n",
    "\n",
    "print(\"Tensorflow\", tf.__version__)\n",
    "print(\"Python\", sys.version)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Post Training Quantization\n",
    "\n",
    "We can take a pre-trained Tensorflow graph and convert it into quantized TFlite graph. To do that, we will only need two things: the range (min and max values) of weights and activation. The former is easy, since the weights don't change after training, we can therefore work out the range from frozen graph. To recall basic of neural network, activation is the output of a layer and the range depend on both the weights and the inputs. Thus, we can't get the range of activation directly from frozen graph. Unless, the activation range is fixed by design, e.g. the non-linearity tf.nn.relu6 cap the value between 0.0 and 6.0. This is also the reason why Google uses relu6 instead of relu in quantized Mobilenet as the latter has no upperbound. However, even if the range i sknown, it is still less than ideal we may lose some granularity. To give a concrete example, for range of (0.0, 6.0), the granularity of 8 bit quantization is 6.0/255 = 0.023529412, meaning the number is discretized into multiple of 0.023529412. If the actual range with real data is within 0.0 and 1.0, then the granularity improved six fold to 1.0/255 = 0.003921569. \n",
    "\n",
    "You can do post training quantization using either Python APIs or command line. You can use different graph format e.g. frozen graph, saved model, from session etc which I find it quite confusing. I'll go through the fundamental using a mixture of format (Python/command line/graph format) and you can refer to the many online examples provided by Tensorflow best suited for your project. Instead of treating them like black box, we'll go through the examples from bottom-up, starting by quantizing a convolution layer using Toco converter (just another fancy acronym came up by Google engineers)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Froze 2 variables.\n",
      "INFO:tensorflow:Converted 2 variables to const ops.\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "\n",
    "# Create a simple network\n",
    "def simple_network(input):\n",
    "\n",
    "    x = tf.layers.conv2d(input, filters=32, kernel_size=3)\n",
    "    x = tf.nn.relu(x)\n",
    "    return x\n",
    "\n",
    "input_dim = [1, 224, 224, 3]\n",
    "input = tf.placeholder(tf.float32, input_dim)\n",
    "output = simple_network(input)\n",
    "\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    \n",
    "    # pass in the graph's input and output\n",
    "    converter = tf.contrib.lite.TocoConverter.from_session(sess, [input], [output])\n",
    "    # set inference type to uint8\n",
    "    converter.inference_type = tf.contrib.lite.constants.QUANTIZED_UINT8\n",
    "    # set the activation range. Try comment out this line and the conversion will fail\n",
    "    converter.default_ranges_stats = (0., 6.)\n",
    "    \n",
    "    input_mean = 128\n",
    "    input_stddev = 128\n",
    "    input_arrays = converter.get_input_arrays()\n",
    "    # the input mean and standard deviation is needed to work out the scale and offset\n",
    "    # to de-quantize the input, since we assume input is quantized. Therefore, we can use\n",
    "    # image's raw RGB uint8 as input directly.\n",
    "    converter.quantized_input_stats = {input_arrays[0] : (input_mean, input_stddev)}  # mean, std_dev\n",
    "    \n",
    "    # now convert\n",
    "    tflite_model = converter.convert()\n",
    "    \n",
    "    # now you can save the quantized model\n",
    "    save_path = \"models/practice\"\n",
    "    if not os.path.exists(save_path):\n",
    "        os.makedirs(save_path)\n",
    "        \n",
    "    open(os.path.join(save_path, \"simple_model.tflite\"), \"wb\").write(tflite_model)\n",
    "    \n",
    "    # you can start using it right now\n",
    "    # load it into interpreter and you can use it like in Tutorial 1\n",
    "    interpreter = tf.contrib.lite.Interpreter(model_content=tflite_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantization Aware Training\n",
    "\n",
    "Converting a \"normal\" frozen graph will normally result in loss in accuracy for two reasons:\n",
    "1. Un-optimized activation range\n",
    "2. Quantization errors\n",
    "\n",
    "Among the two, the quantizaztion error is the greater contributor to accuracy loss. During training, the neural network use full precision, say for example, a value of 1.2345, but during inference, it can see the dequantize value of say 1.1 which is different from what it sees in training. This error will propagate and accumulate as it traverse across the layers which can result in quite some difference at the output.\n",
    "\n",
    "To tackle these 2 problems, we'll need to do quantization aware training whicih insert additionial operations to :\n",
    "1. simulate the quantization effect, and\n",
    "2. measure the activation range\n",
    "\n",
    "Tensorflow has a built-in function to add those operation to your graph. You can look at the website [here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/quantize/README.md)\n",
    "```\n",
    "tf.contrib.quantize.create_training_graph(input_graph=g,\n",
    "                                          quant_delay=2000000)\n",
    "```\n",
    "That's it, just one instruction to prep the graph to be ready for training. The quant_delay (in number of steps) is to delay the fake quantization till the network activation range is more stable. \n",
    "\n",
    "Hooray, jobs done! Hm... yes but if you are curious about what's happening in the black box, then let's dive in for the details. We will build on the simple_network we've just created, and add a few lines to show what actually happens in the transformed graph.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Digging Into Details\n",
    "\n",
    "![sim](images/sim_quant_graph)\n",
    "The graph above is taken from Google's paper [\"Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference\"](https://arxiv.org/abs/1712.05877). It looks straight forward, we only need to implement two fake quantizations, one for weights and one for activation.\n",
    "\n",
    "In simple_network, we use high level function tf.layers to define the convolution layer which takes care of creation and initization of weights and bias. Now since we want to do the quantization on weights and activations ourselves, we'll use lower level APIs tf.get_variable to create the weights and tf.nn to perform the convolution. \n",
    "\n",
    "We'll use tf.fake_quant_with_min_max_vars which is almost identical with tf.fake_quant_with_min_max_args introduced earlier. It is not clear to me what their difference except the input min and max are made mandatory in tf.fake_quant_with_min_max_vars. From the converted graph and source code, I can see only tf.fake_quant_with_min_max_vars therefore we'll use that from now on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "\n",
    "def less_simple_network(input):\n",
    "    '''\n",
    "        Create weights\n",
    "    '''\n",
    "    # [filter_height, filter_width, in_channels, out_channels]\n",
    "    w_dim = [3, 3, 3, 32] \n",
    "    \n",
    "    w = tf.get_variable(\"weight\", \n",
    "                        shape=w_dim,\n",
    "                        initializer=tf.contrib.layers.xavier_initializer())\n",
    "        \n",
    "    '''\n",
    "        Fake quantizer weights\n",
    "    '''\n",
    "    w_min = tf.reduce_min(w)\n",
    "    w_max = tf.reduce_max(w)\n",
    "    \n",
    "    w_fake_quant = tf.fake_quant_with_min_max_vars(w, \n",
    "                    min=w_min, \n",
    "                    max=w_max, \n",
    "                    narrow_range=True, # will be explained below\n",
    "                    name=\"quant_weights\")\n",
    "\n",
    "    '''\n",
    "        Create bias but don't quantize it for reasons to be explained later\n",
    "    '''\n",
    "    bias = tf.get_variable(\"bias\", \n",
    "                        shape=[32],\n",
    "                        initializer=tf.initializers.zeros)\n",
    "    '''\n",
    "        Perform convolution and relu\n",
    "    '''\n",
    "    strides = 1\n",
    "    out = tf.nn.conv2d(input, w_fake_quant, [1, strides, strides, 1], padding='SAME')\n",
    "    out = tf.nn.bias_add(out, bias, name='bias')    \n",
    "    out = tf.nn.relu6(out)\n",
    "    \n",
    "    '''\n",
    "        Fake quantize activation\n",
    "    '''\n",
    "    out_fake_quant = tf.fake_quant_with_min_max_vars(out, \n",
    "                    min=0.0, \n",
    "                    max=6.0, \n",
    "                    narrow_range=False,\n",
    "                    name=\"act_weights\")\n",
    "    \n",
    "    return out_fake_quant\n",
    "\n",
    "input_dim = [1, 224, 224, 3]\n",
    "input2 = tf.placeholder(tf.float32, input_dim)\n",
    "output2 = less_simple_network(input2)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alright, I admit now the code looks much longer because of the lower level APIs but the idea is not that complicated as shown in the diagram above, just add two quantizaton nodes. Now, if we comment out the converter option that set the default_ranges_stats, it will no longer crash because the graph already contains all the range information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Froze 2 variables.\n",
      "INFO:tensorflow:Converted 2 variables to const ops.\n"
     ]
    }
   ],
   "source": [
    "with tf.Session() as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    \n",
    "    # pass in the graph's input and output\n",
    "    converter = tf.contrib.lite.TocoConverter.from_session(sess, [input2], [output2])\n",
    "    # set inference type to uint8\n",
    "    converter.inference_type = tf.contrib.lite.constants.QUANTIZED_UINT8\n",
    "    # Conversion won't crash now because range is set inside fake_quant code\n",
    "    #converter.default_ranges_stats = (0., 6.)\n",
    "    \n",
    "    input_mean = 128\n",
    "    input_stddev = 128\n",
    "    input_arrays = converter.get_input_arrays()\n",
    "    # the input mean and standard deviation is needed to work out the scale and offset\n",
    "    # to de-quantize the input, since we assume input is quantized. Therefore, we can use\n",
    "    # image's raw RGB uint8 as input directly.\n",
    "    converter.quantized_input_stats = {input_arrays[0] : (input_mean, input_stddev)}  # mean, std_dev\n",
    "    \n",
    "    # now convert\n",
    "    tflite_model = converter.convert()\n",
    "    \n",
    "    # now you can save the quantized model\n",
    "    save_path = \"models/practice\"\n",
    "    if not os.path.exists(save_path):\n",
    "        os.makedirs(save_path)\n",
    "        \n",
    "    open(os.path.join(save_path, \"less_simple_model.tflite\"), \"wb\").write(tflite_model)\n",
    "    \n",
    "    # you can start using it right now\n",
    "    # load it into interpreter and you can use it like in Tutorial 1\n",
    "    interpreter = tf.contrib.lite.Interpreter(model_content=tflite_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quantization Range\n",
    "#### Weights\n",
    "As explained earlier, the range of weight's real number is simply the min and max values. However, on the quantize range, there is an option in the function known as narrow_range. We mentioned in earlier tutorials, the range is $2^N-1$ which is 255 for 8 bits so the quantized signed a bit range from -128 to 127. When narrow_range is set, then 254 is used instead hence -127 to 127. I initially thought it was because of making sure value 0 is in the center with equal number of positive and negative values. It turns out that Google want to make use of a ARM processor instruction that require smaller number range to prevent overflow in calculating matrix multiplication. So, that's it. On the other hand, activation uses the full 'wide range'.\n",
    "\n",
    "#### Activation\n",
    "In the example above, I hard code min and max to be 0.0 and 6.0. However, those numbers should be dynamic determined by statistics during training. In practice, an exponential moving average values are used to prevent sudden change from batch to batch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Command Lines & Graph Visualization\n",
    "\n",
    "Before the end of the tutorial, I thought I should introduce the command line to do quantization too. It has a powerful feature that the Python APIs don't have - creating graph visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "converted_model.tflite\r\n",
      "graphviz\r\n",
      "mobilenet_v1_1.0_224_quant.ckpt.data-00000-of-00001\r\n",
      "mobilenet_v1_1.0_224_quant.ckpt.index\r\n",
      "mobilenet_v1_1.0_224_quant.ckpt.meta\r\n",
      "mobilenet_v1_1.0_224_quant_eval.pbtxt\r\n",
      "mobilenet_v1_1.0_224_quant_frozen.pb\r\n",
      "mobilenet_v1_1.0_224_quant_info.txt\r\n",
      "mobilenet_v1_1.0_224_quant.tflite\r\n",
      "mobilenet_v1_1.0_224_quant.tgz\r\n"
     ]
    }
   ],
   "source": [
    "# I run this on Linux, it if doesn't work on your platform, \n",
    "# skip this step and look at directory using your usual ways\n",
    "! ls models/mobilenet_v1/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: mobilenet_v1_1.0_224_quant\r\n",
      "Input: input\r\n",
      "Output: MobilenetV1/Predictions/Reshape_1\r\n"
     ]
    }
   ],
   "source": [
    "# Look at the info file to find out about the input and output nodes\n",
    "! cat models/mobilenet_v1/mobilenet_v1_1.0_224_quant_info.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mkdir: cannot create directory ‘models/mobilenet_v1/graphviz’: File exists\n",
      "2018-11-25 00:15:55.098293: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA\n",
      "2018-11-25 00:15:55.171286: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:897] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2018-11-25 00:15:55.171661: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1405] Found device 0 with properties: \n",
      "name: GeForce GTX 1050 major: 6 minor: 1 memoryClockRate(GHz): 1.493\n",
      "pciBusID: 0000:01:00.0\n",
      "totalMemory: 3.95GiB freeMemory: 3.12GiB\n",
      "2018-11-25 00:15:55.171695: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1484] Adding visible gpu devices: 0\n",
      "2018-11-25 00:15:55.373946: I tensorflow/core/common_runtime/gpu/gpu_device.cc:965] Device interconnect StreamExecutor with strength 1 edge matrix:\n",
      "2018-11-25 00:15:55.373996: I tensorflow/core/common_runtime/gpu/gpu_device.cc:971]      0 \n",
      "2018-11-25 00:15:55.374003: I tensorflow/core/common_runtime/gpu/gpu_device.cc:984] 0:   N \n",
      "2018-11-25 00:15:55.374169: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1097] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 2836 MB memory) -> physical GPU (device: 0, name: GeForce GTX 1050, pci bus id: 0000:01:00.0, compute capability: 6.1)\n"
     ]
    }
   ],
   "source": [
    "# convert the graph\n",
    "# if you get CUDA memory error, restart the jupyter kernel and skip the Python experiments above\n",
    "! mkdir models/mobilenet_v1/graphviz\n",
    "! tflite_convert \\\n",
    "--graph_def_file=models/mobilenet_v1/mobilenet_v1_1.0_224_quant_frozen.pb \\\n",
    "--output_file=models/mobilenet_v1/converted_model.tflite \\\n",
    "--input_arrays=input \\\n",
    "--output_arrays=MobilenetV1/Predictions/Reshape_1 \\\n",
    "--dump_graphviz_dir=models/mobilenet_v1/graphviz \n",
    "# comment out the above line and below if you don't want to look at the graph\n",
    "# in Linux, install graphviz by doing \"sudo apt install graphvizs\"\n",
    "\n",
    "!dot -Tpdf -O models/mobilenet_v1/graphviz/*.dot "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You'll find three pdf in the folder mobilenet_v1/graphviz. One is the graph at import time which is the full Tensorflow graph. It also give you a glimpse on how complex a Tensorflow graph can be, showing every single operation as a node in graph (e.g. add, read, assign). \n",
    "\n",
    "![](images/import_graph.png)\n",
    "\n",
    "During the conversion, the graph will be transformed and simplified before quantization most notably the batchnormalization will be folded into convolution weight and bias and therefore you won't see it the quantized graph.\n",
    "\n",
    "![](images/transformed_graph.png)\n",
    "\n",
    "There is a intermediate graph with transient allocation information which I have no single clue what that is. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's Next?\n",
    "\n",
    "If you want to learn more about the graph conversion Python APIs and command lines, you can find them here for [Python APIs](https://www.tensorflow.org/lite/convert/python_api) and [command lines](https://www.tensorflow.org/lite/convert/cmdline_examples). I believe you have mastered the fundamental of quantizing neural network and is ready to read the original academic paper [\"Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference\"](https://arxiv.org/abs/1712.05877)."
   ]
  }
 ],
 "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
}