Quant_API.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Quantization schemes\n",
    "<img src=\"./img/q_scheme.png\" width=\"600\" />\n",
    "\n",
    "Two sets of schemes:\n",
    "* Symmetric\n",
    "* Affine\n",
    "\n",
    "And\n",
    "\n",
    "* Per-channel\n",
    "* Per-Tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Per-Channel and Per-Tensor\n",
    "\n",
    "<img src=\"./img/per_t_c.png\" width=\"600\" />"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor([\n",
    "    [0.5827, 0.8619], \n",
    "    [0.3827, -0.1982], \n",
    "    [-0.8213, 0.6351]])\n",
    "\n",
    "print(x.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# per-tensor\n",
    "\n",
    "scale = torch.tensor(1e-2)\n",
    "zero_pt = torch.tensor(0)\n",
    "\n",
    "xq = torch.quantize_per_tensor(x, scale, zero_pt, dtype=torch.qint8)\n",
    "print(xq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# per-channel\n",
    "\n",
    "channel_axis = 0\n",
    "scale = torch.tensor([1e-2, 1e-3, 5e-2])\n",
    "zero_pt = torch.zeros(3)\n",
    "\n",
    "xq = torch.quantize_per_channel(x, scale, zero_pt, dtype=torch.qint8, axis=0)\n",
    "print(xq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Symmetric and Affine\n",
    "\n",
    "Symmetric\n",
    "* Input range is calculated symmetrically around 0\n",
    "* Good for quantizing weights\n",
    "* Wasteful for quantizing activations - why?\n",
    "\n",
    "Affine \n",
    "* Clips the input tightly \n",
    "\n",
    "\n",
    "<img src=\"./img/affine-symmetric.png\" width=\"600\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Observers\n",
    "\n",
    "<img src=\"./img/observer.png\" width=\"600\" />"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.ao.quantization.observer import MovingAverageMinMaxObserver, HistogramObserver, MovingAveragePerChannelMinMaxObserver\n",
    "\n",
    "size = (3,4)\n",
    "normal = torch.distributions.normal.Normal(0,1)\n",
    "input = [normal.sample(size) for _ in range(3)]\n",
    "\n",
    "observers = [\n",
    "    MovingAverageMinMaxObserver(qscheme=torch.per_tensor_affine), \n",
    "    HistogramObserver(), \n",
    "    MovingAveragePerChannelMinMaxObserver(qscheme=torch.per_channel_symmetric)\n",
    "    ]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for obs in observers:\n",
    "  for x in input: \n",
    "      obs(x) \n",
    "  print(obs.__class__.__name__, obs.calculate_qparams())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### QConfig\n",
    "\n",
    "* High-level abstraction wrapping these knobs in one object\n",
    "* Allows separate configuration for activation and weights of a layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.ao.quantization.observer import MovingAverageMinMaxObserver, MovingAveragePerChannelMinMaxObserver\n",
    "from torch.ao.quantization.qconfig import QConfig\n",
    "\n",
    "my_qconfig = QConfig(\n",
    "  activation=MovingAverageMinMaxObserver.with_args(\n",
    "      qscheme=torch.per_tensor_affine,\n",
    "      dtype=torch.quint8),\n",
    "  weight=MovingAveragePerChannelMinMaxObserver.with_args(\n",
    "      qscheme=torch.per_channel_symmetric)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Default QConfigs out of the box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.quantization.qconfig.default_per_channel_qconfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(torch.quantization.qconfig.default_dynamic_qconfig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(torch.quantization.qconfig.per_channel_dynamic_qconfig)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "5b2c14c5f2a3b21e6c2412c8196f5145870350e81c0b737cae3e5c60eb1e1eac"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('pytorch_p38': conda)",
   "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.8.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}