tutorials/quantization_tutorial.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "608c2fc3",
   "metadata": {},
   "source": [
    "# Basics of Reduced Precision and Quantization\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ac31e8e4",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrates the development of quantized models for 4-bit inferencing using our [model optimization library](https://pypi.org/project/fms-model-optimizer/). [FMS Model Optimizer](https://github.com/foundation-model-stack/fms-model-optimizer/) is a Python framework for the development of reduced precision neural network models which provides state-of-the-art quantization techniques, together with automated tools to apply these techniques in Pytorch environments for Quantization-Aware-Training (QAT) of popular deep learning workloads. The resulting low-precision models can be deployed on GPUs or other accelerators.\n",
    "\n",
    "We will demonstrate the following:\n",
    "- How input data can be quantized\n",
    "- How quantization can be applied to a convolution layer\n",
    "- How to automate the quantization process\n",
    "- How a quantized convolution layer performs on a typical image\n",
    "\n",
    "FMS Model Optimizer can be applied across a variety of computer vision and natural language processing tasks to speed up inferencing,  reduce power requirements, and reduce model size, while maintaining comparable model accuracy."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bedbc959",
   "metadata": {},
   "source": [
    "## Table of Contents\n",
    "\n",
    "* <a href=\"#fms_mo_quantizer\">Step 1. Quantize a normal data distribution</a>\n",
    "    * <a href=\"#fms_mo_import\">Import code libraries</a>\n",
    "    * <a href=\"#geninput\">Generate input data</a>\n",
    "    * <a href=\"#clip\">Clip input data </a>\n",
    "    * <a href=\"#quant\">Scale, shift, and quantize data</a>\n",
    "    * <a href=\"#dequant\">Dequantize data </a>\n",
    "* <a href=\"#conv\">Step 2. Quantize a convolution layer</a>\n",
    "    * <a href=\"#3p2\">Generate input data</a>\n",
    "    * <a href=\"#3p3\">Quantize input data</a>\n",
    "    * <a href=\"#3p4\">Create a single layer convolution network</a>\n",
    "    * <a href=\"#3p5\">Generate weights and bias</a>\n",
    "    * <a href=\"#3p6\">Quantize weights</a>\n",
    "    * <a href=\"#3p7\">Feed quantized data, weights, and bias into convolution layer</a>\n",
    "* <a href=\"#fms_mo\">Step 3. Use FMS Model Optimizer to automate quantization</a>\n",
    "* <a href=\"#fms_mo_visual\">Step 4. Try a convolution layer on a quantized image</a>\n",
    "* <a href=\"#fms_mo_conclusion\">Conclusion</a>\n",
    "* <a href=\"#fms_mo_learn\">Learn more</a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7f32d95b",
   "metadata": {},
   "source": [
    "<a id=\"`fms_mo`_quantizer\"></a>\n",
    "    \n",
    "## Step 1. Quantize a normal data distribution\n",
    "\n",
    "In this section we show how quantization works using a randomly generated normal distribution of input data. We will feed the input data to a quantizer and show the quantized output.\n",
    "\n",
    "The quantizer can be summarized by the following equations: <br><br>\n",
    "<font size=4>\n",
    "$y_{int} = \\lfloor \\frac{clamp(y, \\alpha_l, \\alpha_u) - zp}{\\Delta} \\rceil$<br>\n",
    "$y_q = y_{int} \\times \\Delta + zp$\n",
    "</font>\n",
    "\n",
    "where:\n",
    "<br>\n",
    "$y$ = input data<br>\n",
    "$y_{int}$ = $y$ transformed and scaled into integer space, e.g. [0, 1, 2, ...]<br>\n",
    "$y_q$ = quantized output from the quantizer<br>\n",
    "$\\alpha_l$, $\\alpha_u$ = lower and upper clip bound for $y$, $\\in \\mathbb{R}$, e.g. [-1.5, 1.5]. Often referred to as clip_vals<br> \n",
    "$\\Delta$ = stepsize between each quantized bin<br>\n",
    "$zp$ = zero point: consider it as clip_min ($\\alpha_l$) in this scenario <br><br>\n",
    "\n",
    "\n",
    "These two equations are commonly known as quantization and dequantization, respectively. \n",
    "We will walk through these equations step-by-step in the cells below."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "127668fe",
   "metadata": {},
   "source": [
    "<a id=\"`fms_mo`_import\"></a>\n",
    "\n",
    "### Import code libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2af65016",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install fms-model-optimizer\n",
    "! pip install wget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7ba7fb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from fms_mo import qconfig_init, qmodel_prep\n",
    "\n",
    "import argparse\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "import pandas as pd\n",
    "import os\n",
    "from os import path"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e9a95377",
   "metadata": {},
   "source": [
    "\n",
    "<a id=\"geninput\"></a>\n",
    "\n",
    "### Generate input data\n",
    "\n",
    "For simplicity, we generate a normal distribution of input data, with the mean set to 0 and standard deviation set to 1. A sample size of 1 million is chosen.\n",
    "\n",
    "The histogram (distribution plot) is shown below, with the y-axis set as a distribution density.\n",
    "\n",
    "For the purpose of this tutorial, we will focus on only the forward pass as shown by the equation. The backward pass indicated by $$ \\frac{dL}{dy_q}, \\frac{dy_q}{dy}, \\frac{dy_q}{d\\alpha_l}, \\frac{dy_q}{d\\alpha_u} $$ will not be covered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10ac1169",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate normal distribution of data with mean=0, std=1, sample size 1e6\n",
    "# Using numpy (much faster than torch on CPU)\n",
    "raw_data = np.random.normal(0, 1, int(1e6))\n",
    "\n",
    "# Plotting the histogram.\n",
    "plt.figure(figsize=(16, 10))\n",
    "plt.hist(raw_data, density=True, bins=128, alpha=0.8, label='y')\n",
    "#plt.legend(loc='upper right')\n",
    "plt.xlabel(\"Data\")\n",
    "plt.ylabel(\"density\")\n",
    "plt.title(\"Normal distribution of data with mean=0, std=1, sample size 1e6\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8cb1fd4d",
   "metadata": {},
   "source": [
    "<a id=\"clip\"></a>\n",
    "\n",
    "### Clip input data\n",
    "\n",
    "Quantization of a tensor means that we can only use a limited number of distinct values (16 in the case of 4-bit precision) to represent all the numbers in the tensor. For 4-bit precision, we will need to:\n",
    "- determine the range we want to represent, i.e. $[ -\\infty, \\infty] => [ \\alpha_l , \\alpha_u]$, which means anything above $\\alpha_u$ or below $\\alpha_l$ will become $\\alpha_u$ and $\\alpha_l$ respectively.\n",
    "- uniformly divide $[\\alpha_l , \\alpha_u]$ into 16 bins and round the numbers to the nearest bin.\n",
    "\n",
    "Ideally, we want to represent the entire range of the original tensor. But because we have limited bins available, we must clip off extreme values. \n",
    "\n",
    "$[\\alpha_l , \\alpha_u]$ can be determined in many ways, including gradual learning through training, or minimizing the quantization error of the tensor. Quantization error can be defined as $$MeanSquaredError(y, y_q) = \\sum_N \\frac{1}{N} {(y - y_q)^2 }$$ , although there could be other definitions. \n",
    "\n",
    "Here we arbitrarily select a clip min and clip max value of -2.5 and 2.5 and perform the following: <br>\n",
    "$$ clamp(y,\\alpha_l,\\alpha_u) $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d894d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define clipped values for upper and lower bound.\n",
    "clip_min, clip_max = -2.5, 2.5\n",
    "clipped_data = np.clip(raw_data, clip_min, clip_max)#clip_data(raw_data, clip_min, clip_max)\n",
    "\n",
    "print( \"min/max of the original tensor\", np.min(raw_data), np.max(raw_data) )\n",
    "print( \"min/max of the clipped tensor\", np.min(clipped_data), np.max(clipped_data) )\n",
    "\n",
    "print(\"MSE(raw_data, clipped_data)\", np.mean( (raw_data-clipped_data)**2 ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b77562e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show the first 5 clipped numbers and their original values.\n",
    "isClipped=np.logical_or(raw_data>clip_max, raw_data<clip_min)\n",
    "idx_clipped_elements=np.where( isClipped )[0]\n",
    "pd.DataFrame( \n",
    "                {'idx':idx_clipped_elements[:5], \n",
    "                'raw': raw_data[ idx_clipped_elements[:5] ],\n",
    "                'clipped': clipped_data[idx_clipped_elements[:5]] }\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49233807",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the distribution and the clipped data to visualize\n",
    "\n",
    "plt.figure(figsize=(16, 10))\n",
    "plt.hist(raw_data,     density=True, bins=64, label=\"y (raw values)\", histtype='step', linewidth=3.5),\n",
    "plt.hist(clipped_data, density=True, bins=64, color=['#33b1ff'], alpha=0.8,label=\"y_clamp (clipped edges)\"), \n",
    "plt.legend(fancybox=True, ncol=2)\n",
    "plt.xlabel(\"Data\")\n",
    "plt.ylabel(\"density\")\n",
    "plt.title(\"Raw Data and Clipped Data\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ac13dfd7",
   "metadata": {},
   "source": [
    "From the results above, we can see that we've successfully clamped the data. \n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0989cffb",
   "metadata": {},
   "source": [
    "<a id=\"quant\"></a>\n",
    "\n",
    "### Scale, shift, and quantize data\n",
    "\n",
    "Here we choose to use 4-bit integer for this quantization, with zp = clip_min. \n",
    "\n",
    "Our next step is to transform the data from the range [-2.5, 2.5] to the range [0, 15] and round the values to the nearest integer. Apply: <br>\n",
    "$$y_{int} = \\lfloor \\frac{clamp(y, \\alpha_l, \\alpha_u) - zp}{\\Delta} \\rceil$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6807c54d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set bit size for quantization\n",
    "n_bit = 4\n",
    "zp = clip_min\n",
    "stepsize = (clip_max - zp) / (2 ** n_bit -1)\n",
    "y_scaled = (clipped_data - clip_min) / stepsize\n",
    "y_int    = np.round(y_scaled)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1205162c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(16, 10))\n",
    "plt.hist(raw_data, density=True, bins=64, alpha=0.8,label=\"y (raw values)\", histtype='step', linewidth=3.5)\n",
    "plt.hist(y_scaled, density=True,  bins=64, color=['#33b1ff'], alpha=0.6,label=\"scale+shift\")\n",
    "plt.hist(y_int,    density=True,  bins=64, color=['#007d79'],alpha=0.8,label=\"quantize\")\n",
    "plt.legend(loc='upper left', fancybox=True, ncol=3)\n",
    "plt.xlabel(\"Data\")\n",
    "plt.ylabel(\"density\")\n",
    "#plt.yscale('log')\n",
    "plt.title(\"Raw Data and Shifted Data\")\n",
    "plt.show()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0c9551fd",
   "metadata": {},
   "source": [
    "The plot above shows that we can represent the data as integers by clipping, shifting, scaling, and quantizing the data."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "01f817a0",
   "metadata": {},
   "source": [
    "<a id=\"dequant\"></a>\n",
    "\n",
    "### Dequantize data\n",
    "\n",
    "The last step is to dequantize the quantized data $y_{int}$ back to the range [-2.5, 2.5] so that it overlays the original distribution. <br>\n",
    "<font size=4>\n",
    "$$y_q = y_{int} \\times \\Delta + zp$$\n",
    "</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cca70648",
   "metadata": {},
   "outputs": [],
   "source": [
    "yq = y_int * stepsize + zp\n",
    "\n",
    "plt.figure(figsize=(16, 10))\n",
    "plt.hist(raw_data, density=True, bins=64, label=\"original y\", histtype='step', linewidth=2.5)#alpha=0.8,\n",
    "plt.hist(yq,       density=True, color=['#33b1ff'], bins=64, label=\"quantized y\")#alpha=0.7,\n",
    "plt.legend(fancybox=True, ncol=2)\n",
    "plt.xlabel(\"Data\")\n",
    "plt.ylabel(\"density\")\n",
    "plt.title(\"Raw Data and Quantized Data\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e794bc37",
   "metadata": {},
   "source": [
    "### An example of symmetric vs asymmetric quantization "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d592edd",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplots(3,1, figsize=(16, 12), sharex=True)\n",
    "\n",
    "arstyle=dict(facecolor='C1',alpha=0.5, shrink=0.05)\n",
    "\n",
    "n_bit = 4\n",
    "clip_min, clip_max = -2.5, 2.5\n",
    "asym_raw_data = np.abs(raw_data)\n",
    "for i, (raw_i, lbl_i) in enumerate([(raw_data, 'Case 1: sym data, sym Q'), \n",
    "                                    (asym_raw_data, 'Case 2: asym data, asym Q'),\n",
    "                                    (asym_raw_data, 'Case 3: asym data sym Q') ]):\n",
    "    if 'asym Q' in lbl_i:\n",
    "        # asym quantization for range [0, clip_max]\n",
    "        clip_min_i = np.min(raw_i)\n",
    "        nbins = 2**n_bit -1\n",
    "        scale = (clip_max - clip_min_i)/nbins\n",
    "        zp = np.round(-clip_min_i/scale)\n",
    "    else:\n",
    "        # sym quantization\n",
    "        clip_min_i = -max(clip_max, np.abs(clip_min))\n",
    "        nbins = 2**n_bit -2\n",
    "        scale = (clip_max - clip_min_i)/nbins\n",
    "        zp = 0\n",
    "\n",
    "    # here we could use one of the 2 commonly used formulas\n",
    "    # 1. y_q = round( (clamp(x) - zp)/scale )*scale + zp\n",
    "    # 2. y_q = (clamp(round(x/scale + zero_point), quant_min, quant_max) - zero_point) * scale\n",
    "    y_int_i = np.round( (np.clip(raw_i, clip_min, clip_max) - zp)/scale )\n",
    "    yq_i = y_int_i*scale + zp\n",
    "    max_bin_i = np.round( (clip_max-zp)/scale)*scale + zp\n",
    "\n",
    "    plt.subplot(311+i)\n",
    "    plt.hist(raw_i, density=False, bins=64, label=\"original y\", histtype='step', linewidth=2.5)\n",
    "    plt.hist(yq_i,  density=False, color=['#33b1ff'], bins=64, label='y_q')\n",
    "    plt.legend(fancybox=True, ncol=2, fontsize=14)\n",
    "\n",
    "    plt.ylabel(\"Count\")\n",
    "    plt.annotate('upper clip bound', xy=(max_bin_i, 0), xytext=(max_bin_i, 1e5), arrowprops=arstyle)    \n",
    "    plt.annotate('lower clip bound', xy=(clip_min_i, 0), xytext=(clip_min_i, 1e5), arrowprops=arstyle)    \n",
    "    plt.title(lbl_i)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "95c539bd",
   "metadata": {},
   "source": [
    "Now we will wrap the above steps into a \"simple quantizer\" so that we can easily reuse it later (using torch instead of numpy)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed62fed0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simpleQuantizer(input, n_bit, clip_min, clip_max):\n",
    "    zp = clip_min\n",
    "    stepsize = (clip_max - zp) / (2 ** n_bit -1)\n",
    "\n",
    "    y_scaled = (torch.clamp(input, clip_min, clip_max) - clip_min) / stepsize\n",
    "    y_int    = torch.round(y_scaled)\n",
    "    return y_int * stepsize + zp\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3562f6ec",
   "metadata": {},
   "source": [
    "<a id=\"conv\"></a>\n",
    "\n",
    "## Step 2. Quantize a convolution layer\n",
    "\n",
    "In this section, we show how to manually quantize a Convolution layer, i.e. quantizing the input data and weights, and then feed them into a convolution computation. \n",
    "\n",
    "**Note:**\n",
    "1. The quantizers here use different clip values and can be different type of quantizers if needed. \n",
    "2. In practice, \"bias\" in convolution layer usually doesn't get quantized. Simply because the computation is much lower compare to matmul and the risk of losing accuracy is very high."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "44cd359c",
   "metadata": {},
   "source": [
    "<a id=\"3p2\"></a>\n",
    "\n",
    "### Generate input data\n",
    "\n",
    "Similar to Step 1, the input data is a randomly generated normal distribution. We generate 1 input sample with 3 channels, 32 pixel width, and 32 pixel height."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b438226",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Channel, Width, Height\n",
    "C, H, W = 3, 32, 32\n",
    "N = 1\n",
    "\n",
    "# Generate 1 sample\n",
    "input = torch.randn(N,C,H,W)\n",
    "\n",
    "print('Input Shape: ', input.shape)\n",
    "print('Number of unique input values: ', input.detach().unique().size()[0])\n",
    "print(f'Expected: {N * C * H * W} (Based on randomly generated values for shape {N} x {C} x {H} x {W})')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4c8da2aa",
   "metadata": {},
   "source": [
    "<a id=\"3p3\"></a>\n",
    "\n",
    "### Quantize input data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d918346a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the max, min clip values and number of bits\n",
    "clip_min, clip_max = -2.5, 2.5\n",
    "n_bit = 4\n",
    "\n",
    "# Quantize the input data\n",
    "input_quant = simpleQuantizer(input, n_bit, clip_min, clip_max)\n",
    "\n",
    "print('Quantized input Shape: ', input_quant.shape)\n",
    "print('Number of unique quantized input values: ', input_quant.detach().unique().size()[0])\n",
    "print(f'Expected: {2 ** n_bit} (Based on 2 ^ {n_bit})')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a15bbc6f",
   "metadata": {},
   "source": [
    "<a id=\"3p4\"></a>\n",
    "\n",
    "### Create a single layer convolution network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5db080c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Network of 1 Convolution Layer\n",
    "class NeuralNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNet, self).__init__()\n",
    "        # (32 (width, height) - 3 (filter size)) / 1 (stride) + 1 = 30 (new width, height)\n",
    "        self.conv = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3, stride=1)\n",
    "\n",
    "    def forward(self, input):\n",
    "        out = self.conv(input)\n",
    "        return out\n",
    "\n",
    "net = NeuralNet()\n",
    "\n",
    "net(input).shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9432c215",
   "metadata": {},
   "source": [
    "<a id=\"3p5\"></a>\n",
    "\n",
    "### Generate weights and bias\n",
    "\n",
    "To simulate the quantization of a pretrained model we set the weights manually to a normal distribution of values. Bias will be set to zeros because we don't plan on using bias."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "860b6033",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the weights for the convolution filter so we know what the values are\n",
    "weight = torch.randn(net.conv.weight.shape)\n",
    "bias = torch.zeros(net.conv.bias.shape)\n",
    "\n",
    "# Replace current conv2d weight with this randomly generated weight (so we know the values)\n",
    "net.conv.weight = torch.nn.Parameter(weight)\n",
    "\n",
    "# ignore bias for now \n",
    "net.conv.bias = torch.nn.Parameter(bias)\n",
    "\n",
    "print('Weight Shape: ', weight.shape)\n",
    "print('Number of unique weight values: ', weight.detach().unique().size()[0])\n",
    "print(f'Expected: {weight.numel()} (Based on randomly generated values for shape {weight.shape[0]} x {weight.shape[1]} x {weight.shape[2]} x {weight.shape[3]})')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f76e8398",
   "metadata": {},
   "source": [
    "<a id=\"3p6\"></a>\n",
    "\n",
    "### Quantize weights\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d3e53d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set variables for quantization\n",
    "\n",
    "# Quantize the weights (similar to input)\n",
    "weight_quant = simpleQuantizer(weight, n_bit, clip_min, clip_max)\n",
    "\n",
    "print('Quantized weight Shape: ', weight_quant.shape)\n",
    "print('Number of unique quantized weight values: ', weight_quant.detach().unique().size()[0])\n",
    "print(f'Expected: {2 ** n_bit} (Based on 2 ^ {n_bit})')\n",
    "print('First Channel of Quantized Weight', weight_quant[0])\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5ec402b7",
   "metadata": {},
   "source": [
    "<a id=\"3p7\"></a>\n",
    "\n",
    "### Feed quantized data, weights, and bias into convolution layer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "149cbe60",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate output y\n",
    "y = net(input)\n",
    "\n",
    "# Generate quantized output y, NOTE, this net is currently using non-quantized weight \n",
    "y_quant = net(input_quant)\n",
    "\n",
    "print('Number of unique output values: ', y.detach().unique().size()[0])\n",
    "print('Expected maximum unique output values: ', y.flatten().size()[0])\n",
    "print('Number of unique quantized output values: ', y_quant.detach().unique().size()[0])\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e62182e6",
   "metadata": {},
   "source": [
    "**Now we plot four cases to determine how well quantization works with convolution:**\n",
    "\n",
    "1. both input and weights are not quantized\n",
    "2. quantized weights with raw input\n",
    "3. raw weights with quantized input\n",
    "4. both weights and input are quantized \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c37a059",
   "metadata": {},
   "outputs": [],
   "source": [
    "def PlotAndCompare(d1, d2, labels, title):\n",
    "    mse = nn.functional.mse_loss(d1, d2, reduction='mean' )\n",
    "    plt.hist( d1.flatten().detach().numpy(), bins=64, alpha = 0.7, density=True, label=labels[0])\n",
    "    plt.hist( d2.flatten().detach().numpy(), bins=64, color=['#33b1ff'], alpha = 0.8, density=True, label=labels[1], histtype='step', linewidth=3.5)\n",
    "    plt.yscale('log')\n",
    "    plt.legend(loc='upper center', bbox_to_anchor=(0.5, -0.1), fancybox=True, ncol=2)\n",
    "    plt.title(f\"{title}, MSE={mse:.3f}\")\n",
    "\n",
    "\n",
    "\n",
    "titles=['inputs', 'weights', 'outputs']\n",
    "isQ = ['not quantized', 'quantized']\n",
    "for i, inp in enumerate([input, input_quant]):\n",
    "    for j, W in enumerate([weight, weight_quant]):\n",
    "        plt.subplots(1,3,figsize=(18,5))\n",
    "        plt.suptitle(f'Case {i*2+j+1}: Input {isQ[i]}, Weight {isQ[j]}', fontsize=20, ha='center', va='bottom')\n",
    "        plt.subplot(131); PlotAndCompare(input,     inp,        ['raw', isQ[i]],  f\"input, {isQ[i]}\")\n",
    "        plt.subplot(132); PlotAndCompare(weight,    W,          ['raw', isQ[j]],  f\"weight, {isQ[j]}\")\n",
    "        net.conv.weight = torch.nn.Parameter(W)\n",
    "        y_quant = net(inp)\n",
    "        plt.subplot(133); PlotAndCompare(y,         y_quant,   ['raw', f'A={isQ[j]}, W={isQ[i]}'], \"conv output\")\n",
    "        plt.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "691d3c35",
   "metadata": {},
   "source": [
    "We see different levels of MSE when we quantize different components of the convolution layer. Note that for a given set of input and weight tensors, MSE depends on the `number of bits`, `clip_min`, and `clip_max` chosen."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d5316dc5",
   "metadata": {},
   "source": [
    "<a id=\"fms_mo\"></a>\n",
    "\n",
    "## Step 3. Use FMS Model Optimizer to automate quantization\n",
    "\n",
    "In this section we show how to reduce manual effort in the quantization process by using our model optimization library to automate the process.\n",
    "\n",
    "For simplicity we will use a 1-layer toy network as an example, but FMS Model Optimizer can handle more complicated networks. \n",
    "\n",
    "As in Step 2, to simulate the quantization of a pretrained model we set the weights manually to a normal distribution of values. Bias will be set to zeros because we don't plan on using bias.\n",
    "\n",
    "We initialize the configuration (dictionary), manually modify the parameters of interest, then run a \"model prep\" to quantize the network automatically. The results will be identical to the output `y_quant` shown in Step 2.\n",
    "\n",
    "The parameters `nbits_w` and `nbits_a` will be used to control the precision of (most of) the modules identified by FMS Model Optimizer that can be quantized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8e365af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a neural network with single convolution for fms_mo demo purposes\n",
    "net_fms_mo = NeuralNet()\n",
    "\n",
    "# Set weights and bias in convolution network (optional)\n",
    "# Explain that this step is used to be consistent with previous steps.\n",
    "net_fms_mo.conv.weight = torch.nn.Parameter(weight)\n",
    "net_fms_mo.conv.bias = torch.nn.Parameter(bias)\n",
    "\n",
    "\n",
    "# Step 1: initialize configuration dict\n",
    "qcfg = qconfig_init()\n",
    "\n",
    "# set bits for quantization (nbits_a needs to be set to quantize input regardless of bias)\n",
    "qcfg['nbits_w'] = 4\n",
    "qcfg['nbits_a'] = 4\n",
    "\n",
    "# just to be consistent with our \"simple Quantizer\" (normally align_zero is True)\n",
    "qcfg['align_zero'] = False\n",
    "\n",
    "# Quantization Mode here means which quantizers we would like to use,\n",
    "# There are many quantizers available in fms_mo, such as PArameterized Clipping acTivation (PACT),\n",
    "# Statstics-Aware Weight Binning (SAWB).\n",
    "qcfg['qw_mode'] = 'pact'\n",
    "qcfg['qa_mode'] = 'pact'\n",
    "\n",
    "# Set weight and input (activation) clip vals\n",
    "qcfg['w_clip_init_valn'], qcfg['w_clip_init_val'] = -2.5, 2.5\n",
    "qcfg['act_clip_init_valn'], qcfg['act_clip_init_val'] = -2.5, 2.5\n",
    "\n",
    "\n",
    "# This parameter is usually False, but for Demo purposes we quantize the first/only layer\n",
    "qcfg['q1stlastconv'] = True\n",
    "\n",
    "\n",
    "if path.exists(\"results\"):\n",
    "    print(\"results folder exists!\")\n",
    "else:\n",
    "    os.makedirs('results')\n",
    "    \n",
    "# Step 2: Prepare the model to convert layer to add Quantizers\n",
    "qmodel_prep(net_fms_mo, input, qcfg, save_fname='./results/temp.pt')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9679f5f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 3: Run network as usual\n",
    "y_quant_fms_mo = net_fms_mo(input)\n",
    "y_quant      = net(input_quant) \n",
    "\n",
    "plt.figure(figsize=(16, 10))\n",
    "PlotAndCompare(y_quant_fms_mo, y_quant, ['fms_mo','manual'],'quantized Conv output by different methods')\n",
    "plt.show()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "96a43a41",
   "metadata": {},
   "source": [
    "<a id=\"`fms_mo`_visual\"></a>\n",
    "\n",
    "## Step 4. Try a convolution layer on a quantized image\n",
    "\n",
    "In this section we pass an image of a lion through a quantizer and convolution layer to observe the performance of the quantizer with convolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2736eec6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, wget\n",
    "IMG_FILE_NAME = 'lion.png'\n",
    "url = 'https://raw.githubusercontent.com/foundation-model-stack/fms-model-optimizer/main/tutorials/images/' + IMG_FILE_NAME\n",
    "\n",
    "if not os.path.isfile(IMG_FILE_NAME):\n",
    "  wget.download(url, out=IMG_FILE_NAME)\n",
    "\n",
    "img = Image.open(IMG_FILE_NAME)\n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d515bfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "convert_tensor = transforms.ToTensor()\n",
    "input_img_tensor = convert_tensor(img)\n",
    "input_img_tensor = input_img_tensor.unsqueeze(0)\n",
    "\n",
    "tensor_img_transform = transforms.ToPILImage()\n",
    "\n",
    "# we used unsqueeze to create batch dimension, i.e. from [C,H,W] to [N,C,H,W]\n",
    "print(input_img_tensor.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a9ca562",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "net_img_non_quantized = NeuralNet()\n",
    "\n",
    "# Generate weights\n",
    "weight = torch.randn(net_img_non_quantized.conv.weight.shape)\n",
    "\n",
    "# Replace current conv2d weight (so we know the values)\n",
    "net_img_non_quantized.conv.weight = torch.nn.Parameter(weight)\n",
    "\n",
    "# Since we are recycling the net_fms_mo from previous section, the weight needs to be replaced\n",
    "weight_quant = simpleQuantizer(weight, n_bit, clip_min, clip_max)\n",
    "net_fms_mo.conv.weight = torch.nn.Parameter(weight_quant)\n",
    "\n",
    "# Generate normal output from filter\n",
    "y_img_tensor = net_img_non_quantized(input_img_tensor)\n",
    "y_img_quant  = net_fms_mo(input_img_tensor)\n",
    "\n",
    "# Transform output to image\n",
    "feature_map       = tensor_img_transform(y_img_tensor[0])\n",
    "feature_map_quant = tensor_img_transform(y_img_quant[0])\n",
    "\n",
    "\n",
    "plt.subplots(3,1,figsize=(16,25))\n",
    "plt.subplot(311)\n",
    "plt.title('Output from non-quantized model', fontsize=20)\n",
    "plt.imshow(feature_map, cmap='RdBu')\n",
    "plt.clim(0,255)\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(312)\n",
    "plt.title('Output from quantized model', fontsize=20)\n",
    "plt.imshow(feature_map_quant, cmap='RdBu')\n",
    "plt.clim(0,255)\n",
    "plt.colorbar()\n",
    "\n",
    "plt.subplot(313)\n",
    "PlotAndCompare(y_img_tensor, y_img_quant, ['raw','quantized'],'Conv output')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "239c5623",
   "metadata": {},
   "source": [
    "<mark style=\"background-color: lightyellow\">\n",
    "We can see that many details in the second image from the quantized model are saturated and lost. But the shape of the lion can still be seen clearly. This implies that if the quantized model is properly trained, the most critical information can be preserved. For example, if we want to perform a classification or object detection, we may be able to achieve the same answer that \"It's a Lion!\" from both images.\n",
    "</mark>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "aee9be20",
   "metadata": {},
   "source": [
    "<a id=\"`fms_mo`_conclusion\"></a>\n",
    "\n",
    "## Conclusion\n",
    "\n",
    "This notebook provided the following demonstrations:\n",
    "\n",
    "- In Step 1, we showed how quantization can be applied manually to a randomly generated normal distribution of input data.\n",
    "- In Step 2, we showed how to apply quantization to a convolution layer.\n",
    "- In Step 3, we showed how to automate the quantization process using FMS Model Optimizer.\n",
    "- In Step 4, we observed the performance of a quantized convolution layer on an image of a lion."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2330b5b6",
   "metadata": {},
   "source": [
    "<a id=\"`fms_mo`_learn\"></a>\n",
    "\n",
    "## Learn more \n",
    "\n",
    "Please see [example scripts](https://github.com/foundation-model-stack/fms-model-optimizer/tree/main/examples) for more practical use of FMS Model Optimizer.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_latest",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}