notebooks/Brevitas_TVMCon2021.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Brevitas TVMCon 2021 tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Brevitas requires Python 3.8+ and PyTorch 1.5.0+ and can be installed from PyPI with `pip install brevitas`. This notebook was originally released for Brevitas 0.7.1.\n",
    "It has since been updated to run with Brevitas 0.8 and PyTorch 1.9.0 and is going to be available at `https://github.com/Xilinx/brevitas/tree/master/notebooks`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fundamentals\n",
    "\n",
    "## QuantLinear layer\n",
    "\n",
    "We start by looking at `brevitas.nn.QuantLinear`, a quantized alternative to `torch.nn.Linear`, and an instance of a `QuantWeightBiasInputOutputLayer`, meaning that it supports quantization of its weight, bias, input and output. Other instances of `QuantWBIOL` are `QuantConv1d`, `QuantConv2d`,  `QuantConvTranspose1d` and `QuantConvTranspose2d`, and they all follow the same principles. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "    def __init__(\n",
       "            self,\n",
       "            in_features: int,\n",
       "            out_features: int,\n",
       "            bias: bool,\n",
       "            weight_quant: Optional[WeightQuantType] = Int8WeightPerTensorFloat,\n",
       "            bias_quant: Optional[BiasQuantType] = None,\n",
       "            input_quant: Optional[ActQuantType] = None,\n",
       "            output_quant: Optional[ActQuantType] = None,\n",
       "            return_quant_tensor: bool = False,\n",
       "            **kwargs) -> None:\n",
       "        Linear.__init__(self, in_features, out_features, bias)\n",
       "        QuantWBIOL.__init__(\n",
       "            self,\n",
       "            weight_quant=weight_quant,\n",
       "            bias_quant=bias_quant,\n",
       "            input_quant=input_quant,\n",
       "            output_quant=output_quant,\n",
       "            return_quant_tensor=return_quant_tensor,\n",
       "            **kwargs)\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import inspect\n",
    "from brevitas.nn import QuantLinear\n",
    "from IPython.display import Markdown, display\n",
    "\n",
    "def pretty_print_source(source):\n",
    "    display(Markdown('```python\\n' + source + '\\n```'))\n",
    "    \n",
    "source = inspect.getsource(QuantLinear.__init__)  \n",
    "pretty_print_source(source)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default `weight_quant=Int8WeightPerTensorFloat`, while `bias_quant`, `input_quant` and `output_quant` are set to `None`, meaning that by default weight quantization is enabled, while everything else is disabled."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Weight quantization \n",
    "\n",
    "### Default weight quantization\n",
    "\n",
    "`weight_quant=Int8WeightPerTensorFloat` means that by default weights are quantized to *8-bit signed integer with a per-tensor floating-point scale factor*, while quantization of bias, input, and output are disabled.\n",
    "\n",
    "We can instantiate a layer and inspect the original weight tensor and the quantized version to see the effect:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original float weight tensor:\n",
      " Parameter containing:\n",
      "tensor([[-0.0053,  0.3793],\n",
      "        [-0.5820, -0.5204],\n",
      "        [-0.2723,  0.1896],\n",
      "        [-0.0140,  0.5607]], requires_grad=True) \n",
      "\n",
      "Quantized weight QuantTensor:\n",
      " QuantTensor(value=tensor([[-0.0046,  0.3803],\n",
      "        [-0.5820, -0.5224],\n",
      "        [-0.2704,  0.1879],\n",
      "        [-0.0137,  0.5591]], grad_fn=<MulBackward0>), scale=tensor(0.0046, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True)) \n",
      "\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "torch.manual_seed(0)\n",
    "\n",
    "quant_linear = QuantLinear(2, 4, bias=True)\n",
    "\n",
    "print(f\"Original float weight tensor:\\n {quant_linear.weight} \\n\")\n",
    "print(f\"Quantized weight QuantTensor:\\n {quant_linear.quant_weight()} \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As it can be noticed, by default the scale factor is differentiable and it's computed based on the maximum absolute value found within the full precision weight tensor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also retrieve the corrisponding integer representation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Quantized Weight integer tensor:\n",
      " tensor([[  -1,   83],\n",
      "        [-127, -114],\n",
      "        [ -59,   41],\n",
      "        [  -3,  122]], dtype=torch.int32)\n"
     ]
    }
   ],
   "source": [
    "print(f\"Quantized Weight integer tensor:\\n {quant_linear.quant_weight().int()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mixing quantized weights and floating-point inputs\n",
    "\n",
    "In any of the layers defined so far we enabled only weight quantization. So by passing in a float input we would get in general an float output, with the linear operation being computed between the unquantized input and the dequantized weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Float input:\n",
      " tensor([[ 1.5410, -0.2934],\n",
      "        [-2.1788,  0.5684],\n",
      "        [-1.0845, -1.3986]])\n",
      "Float output:\n",
      " tensor([[-0.9036, -0.4586,  0.3096, -0.6472],\n",
      "        [ 1.2058,  0.6525, -0.3723,  0.8677],\n",
      "        [ 1.3873,  0.2801, -0.9009,  0.9507]], grad_fn=<MmBackward>)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "float_input = torch.randn(3, 2)\n",
    "quant_linear = QuantLinear(2, 4, bias=False)\n",
    "float_output = quant_linear(float_input)\n",
    "\n",
    "print(f\"Float input:\\n {float_input}\")\n",
    "print(f\"Float output:\\n {float_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general operations involving quantized tensors are always computed through standard torch operators (here `torch.nn.functional.linear` called internally by the module) on the dequantized representation, the so-called fake-quantization approach."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed-point weight quantization\n",
    "\n",
    "Brevitas exposes various pre-made quantizers. For example, for pure fixed point quantization (i.e. restricting the scale to a power of two) we can use `Int8WeightPerTensorFixedPoint`: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weight QuantTensor:\n",
      " QuantTensor(value=tensor([[-0.0078,  0.3828],\n",
      "        [-0.5781, -0.5234],\n",
      "        [-0.2734,  0.1875],\n",
      "        [-0.0156,  0.5625]], grad_fn=<MulBackward0>), scale=tensor(0.0078, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))\n",
      "Weight fix point: 7.0\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "import math\n",
    "from brevitas.quant import Int8WeightPerTensorFixedPoint\n",
    "    \n",
    "quant_linear = QuantLinear(2, 4, weight_quant=Int8WeightPerTensorFixedPoint, bias=False)\n",
    "\n",
    "print(f\"Weight QuantTensor:\\n {quant_linear.quant_weight()}\")\n",
    "print(f\"Weight fix point: {- math.log2(quant_linear.quant_weight().scale)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Binary weight quantization\n",
    "\n",
    "For binary quantization with a constant scale factor (defaults 0.1) we can set the `SignedBinaryWeightPerTensorConst` quantizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weight QuantTensor:\n",
      " QuantTensor(value=tensor([[-0.1000,  0.1000],\n",
      "        [-0.1000, -0.1000],\n",
      "        [-0.1000,  0.1000],\n",
      "        [-0.1000,  0.1000]], grad_fn=<MulBackward0>), scale=tensor(0.1000), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.quant import SignedBinaryWeightPerTensorConst\n",
    "    \n",
    "quant_linear = QuantLinear(2, 4, weight_quant=SignedBinaryWeightPerTensorConst, bias=False)\n",
    "\n",
    "print(f\"Weight QuantTensor:\\n {quant_linear.quant_weight()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sharing a weight quantizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Brevitas also allows to share **instances** (not definitions) of weight quantizers among layers. This is different from applying the same quantization strategy to multiple layers, and it forces them to have the same scale, zero-point, and bit-width. It can be handy in a variety of situations as we will see later, and it works across different quantization strategies. \n",
    "\n",
    "It's designed to operate in an eager-mode friendly way, meaning that the number of layers sharing it doesn't have to be known in advance. The quantizer is re-initialized appropriately any time it's shared to a new layer. \n",
    "\n",
    "For example, we mentioned that the default weight quantizer `Int8WeightPerTensorFloat` computes the scale factor based on the maximum value found within the floating-point weight tensor to quantize. Sharing it among multiple layers means that the quantizer now looks at all the weight tensors that are being quantized to determine the overall maximum value and generate a single scale factor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QuantLinear 1 scale before sharing with QuantLinear 2: 0.0046\n",
      "QuantLinear 2 scale: 0.0053\n",
      "QuantLinear 1 scale after sharing with QuantLinear 2: 0.0053\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "# Define a QuantLinear layer 1\n",
    "quant_linear1 = QuantLinear(2, 4, bias=False)\n",
    "\n",
    "# Keep a pointer to the scale factor of QuantLinear layer 1 weights before sharing\n",
    "quant_linear1_scale_before_sharing = quant_linear1.quant_weight().scale\n",
    "\n",
    "# Define a QuantLinear layer 2 where the weight quantizer is taken from layer 1\n",
    "quant_linear2 = QuantLinear(2, 4, weight_quant=quant_linear1.weight_quant, bias=False)\n",
    "\n",
    "print(f\"QuantLinear 1 scale before sharing with QuantLinear 2: {quant_linear1_scale_before_sharing:.4f}\")\n",
    "print(f\"QuantLinear 2 scale: {quant_linear2.quant_weight().scale:.4f}\")\n",
    "print(f\"QuantLinear 1 scale after sharing with QuantLinear 2: {quant_linear1.quant_weight().scale:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inputs/Outputs/Activations quantization:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's say now we want to quantize the input to `QuantLinear` to generate a quantized output. We can set an appropriate input quantizer like `Int8ActPerTensorFloat`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Float input:\n",
      " tensor([[ 1.5410, -0.2934],\n",
      "        [-2.1788,  0.5684],\n",
      "        [-1.0845, -1.3986]]) \n",
      "\n",
      "Quant output:\n",
      " tensor([[-0.9109, -0.4609,  0.3135, -0.6523],\n",
      "        [ 1.2089,  0.6524, -0.3752,  0.8697],\n",
      "        [ 1.3893,  0.2816, -0.9011,  0.9521]], grad_fn=<MmBackward>)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.quant import Int8ActPerTensorFloat\n",
    "\n",
    "float_input = torch.randn(3, 2)\n",
    "quant_linear = QuantLinear(2, 4, input_quant=Int8ActPerTensorFloat, bias=False)\n",
    "\n",
    "quant_output = quant_linear(float_input)\n",
    "\n",
    "print(f\"Float input:\\n {float_input} \\n\")\n",
    "print(f\"Quant output:\\n {quant_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how by default the output tensor is returned as a standard torch tensor in dequantized format. This is designed to minimize friction with standard torch operators in the default case. To return an output `QuantTensor` we have to set `return_quant_tensor=True`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Quant output:\n",
      " QuantTensor(value=tensor([[-0.9109, -0.4609,  0.3135, -0.6523],\n",
      "        [ 1.2089,  0.6524, -0.3752,  0.8697],\n",
      "        [ 1.3893,  0.2816, -0.9011,  0.9521]], grad_fn=<MmBackward>), scale=tensor([[9.0542e-05]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(17.), signed_t=tensor(True), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.quant import Int8ActPerTensorFloat\n",
    "\n",
    "float_input = torch.randn(3, 2)\n",
    "quant_linear = QuantLinear(2, 4, input_quant=Int8ActPerTensorFloat, bias=False, return_quant_tensor=True)\n",
    "\n",
    "quant_output = quant_linear(float_input)\n",
    "\n",
    "print(f\"Quant output:\\n {quant_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### QuantIdentity layer\n",
    "\n",
    "Alternatively we could have instantiated a `QuantIdentity` layer. By default `QuantIdentity` adopts the `Int8ActPerTensorFloat` quantizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Float input:\n",
      " tensor([[ 1.5410, -0.2934],\n",
      "        [-2.1788,  0.5684],\n",
      "        [-1.0845, -1.3986]]) \n",
      "\n",
      "Quant input:\n",
      " QuantTensor(value=tensor([[ 1.5490, -0.2894],\n",
      "        [-2.1788,  0.5617],\n",
      "        [-1.0894, -1.3958]], grad_fn=<MulBackward0>), scale=tensor(0.0170, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True)) \n",
      "\n",
      "Quant output:\n",
      " QuantTensor(value=tensor([[-0.9109, -0.4609,  0.3135, -0.6523],\n",
      "        [ 1.2089,  0.6524, -0.3752,  0.8697],\n",
      "        [ 1.3893,  0.2816, -0.9011,  0.9521]], grad_fn=<MmBackward>), scale=tensor([[9.0542e-05]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(17.), signed_t=tensor(True), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.nn import QuantIdentity\n",
    "\n",
    "float_input = torch.randn(3, 2)\n",
    "quant_identity = QuantIdentity(return_quant_tensor=True)\n",
    "quant_linear = QuantLinear(2, 4, bias=False, return_quant_tensor=True)\n",
    "\n",
    "quant_input = quant_identity(float_input)\n",
    "quant_output = quant_linear(quant_input)\n",
    "\n",
    "print(f\"Float input:\\n {float_input} \\n\")\n",
    "print(f\"Quant input:\\n {quant_input} \\n\")\n",
    "print(f\"Quant output:\\n {quant_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that Having the input/output quantizer as part of a layer like `QuantLinear` or as a standalone `QuantIdentity` can make a difference at export time with targets like e.g. the standard ONNX opset, where `QLinearMatMul` or `QLinearConv` assume the output quantizer is part of the layer. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### QuantReLU layer\n",
    "\n",
    "For QuantReLU the default quantizer is `Uint8ActPerTensorFloat`. Quantized activations layers like `QuantReLU` by defaults quantize the output, meaning that a relu is computed first, and then the output is quantized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Float input:\n",
      " tensor([[ 1.5410, -0.2934],\n",
      "        [-2.1788,  0.5684],\n",
      "        [-1.0845, -1.3986]]) \n",
      "\n",
      "Quant output:\n",
      " QuantTensor(value=tensor([[1.5410, 0.0000],\n",
      "        [0.0000, 0.5681],\n",
      "        [0.0000, 0.0000]], grad_fn=<MulBackward0>), scale=tensor(0.0060, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(False), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.nn import QuantReLU\n",
    "\n",
    "float_input = torch.randn(3, 2)\n",
    "quant_relu = QuantReLU(return_quant_tensor=True)\n",
    "\n",
    "quant_output = quant_relu(float_input)\n",
    "\n",
    "print(f\"Float input:\\n {float_input} \\n\")\n",
    "print(f\"Quant output:\\n {quant_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Requantizing a tensor\n",
    "\n",
    "As you might expect, we can pass a `QuantTensor` to a quantized activation, and the tensor will be requantized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Float input:\n",
      " tensor([[ 1.5410, -0.2934],\n",
      "        [-2.1788,  0.5684],\n",
      "        [-1.0845, -1.3986]]) \n",
      "\n",
      "Quant output after QuantIdentity:\n",
      " QuantTensor(value=tensor([[ 1.5490, -0.2894],\n",
      "        [-2.1788,  0.5617],\n",
      "        [-1.0894, -1.3958]], grad_fn=<MulBackward0>), scale=tensor(0.0170, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))\n",
      "Quant output after QuantReLU:\n",
      " QuantTensor(value=tensor([[1.5490, 0.0000],\n",
      "        [0.0000, 0.5588],\n",
      "        [0.0000, 0.0000]], grad_fn=<MulBackward0>), scale=tensor(0.0061, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(False), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "float_input = torch.randn(3, 2)\n",
    "quant_identity = QuantIdentity(return_quant_tensor=True)\n",
    "quant_relu = QuantReLU(return_quant_tensor=True)\n",
    "\n",
    "signed_quant_output = quant_identity(float_input)\n",
    "unsigned_quant_output = quant_relu(signed_quant_output)\n",
    "\n",
    "print(f\"Float input:\\n {float_input} \\n\")\n",
    "print(f\"Quant output after QuantIdentity:\\n {signed_quant_output}\")\n",
    "print(f\"Quant output after QuantReLU:\\n {unsigned_quant_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### How is the activation scale determined by default?\n",
    "\n",
    "To minimize user interaction, default quantizers like `Uint8ActPerTensorFloat` and `Int8ActPerTensorFloat` initializes activations' scale by collecting statistics for a number of training steps (by default the 99.999 percentile of the absolute value for 300 steps). \n",
    "\n",
    "This can be seen as an initial calibration step, although by default it happens with quantization already enabled, while calibration typically collects floating-point statistics first, and the enables quantization. These statistics are accumulated in an exponential moving average that at end of the collection phase is used to internally initialize a learned `nn.Parameter`. It's also possible to delay enabling quantization and perform proper calibration, but we will cover that later.\n",
    "\n",
    "During the collection phase the quantizer behaves differently between `train()` and `eval()` mode, similarly to what happens for batch normalization. In `train()` mode, the statistics for that particular batch are returned. In `eval()` mode, the exponential moving average is returned. After the collection phase is over the learned parameter is returned in both execution modes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bias Quantization\n",
    "\n",
    "In many inference toolchains, bias is assumed to be quantized with scale factor equal to `input_scale * weight_scale`, which means that we need a quantized input somehow. A predefined bias quantizer that reflects that assumption is `brevitas.quant.scaled_int.Int16Bias`. If we simply tried to set it to a `QuantLinear` without any sort of input quantization, we would get an error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Input scale required",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32mC:\\Users\\ALESSA~1\\AppData\\Local\\Temp/ipykernel_18920/2660651517.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      6\u001b[0m \u001b[0mquant_linear\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mQuantLinear\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m2\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m4\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mTrue\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbias_quant\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mInt16Bias\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mreturn_quant_tensor\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mTrue\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      7\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 8\u001b[1;33m \u001b[0mquant_output\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mquant_linear\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfloat_input\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32m~\\miniconda3\\envs\\pt190\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m   1049\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[0;32m   1050\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[1;32m-> 1051\u001b[1;33m             \u001b[1;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1052\u001b[0m         \u001b[1;31m# Do not call functions when jit is used\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1053\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mc:\\users\\alessandro\\documenti\\brevitas_tvmcon\\src\\brevitas\\nn\\quant_linear.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, input)\u001b[0m\n\u001b[0;32m     96\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     97\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mTensor\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mQuantTensor\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m->\u001b[0m \u001b[0mUnion\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mTensor\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mQuantTensor\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 98\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward_impl\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     99\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    100\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0minner_forward_impl\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mx\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mquant_weight\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mquant_bias\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mTensor\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mc:\\users\\alessandro\\documenti\\brevitas_tvmcon\\src\\brevitas\\nn\\quant_layer.py\u001b[0m in \u001b[0;36mforward_impl\u001b[1;34m(self, inp)\u001b[0m\n\u001b[0;32m    355\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    356\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbias\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 357\u001b[1;33m             \u001b[0mquant_bias\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbias_quant\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0moutput_scale\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0moutput_bit_width\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    358\u001b[0m             \u001b[1;32mif\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtraining\u001b[0m \u001b[1;32mand\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcache_inference_quant_bias\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    359\u001b[0m                 \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_cached_bias\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_CachedIO\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mquant_bias\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmetadata_only\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mFalse\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\miniconda3\\envs\\pt190\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m   1049\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[0;32m   1050\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[1;32m-> 1051\u001b[1;33m             \u001b[1;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1052\u001b[0m         \u001b[1;31m# Do not call functions when jit is used\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1053\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mc:\\users\\alessandro\\documenti\\brevitas_tvmcon\\src\\brevitas\\proxy\\parameter_quant.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, x, input_scale, input_bit_width)\u001b[0m\n\u001b[0;32m    194\u001b[0m             \u001b[0mimpl\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mexport_handler\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mexport_mode\u001b[0m \u001b[1;32melse\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtensor_quant\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    195\u001b[0m             \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrequires_input_scale\u001b[0m \u001b[1;32mand\u001b[0m \u001b[0minput_scale\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 196\u001b[1;33m                 \u001b[1;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m\"Input scale required\"\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    197\u001b[0m             \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrequires_input_bit_width\u001b[0m \u001b[1;32mand\u001b[0m \u001b[0minput_bit_width\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    198\u001b[0m                 \u001b[1;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m\"Input bit-width required\"\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: Input scale required"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.quant.scaled_int import Int16Bias\n",
    "\n",
    "float_input = torch.randn(3, 2)\n",
    "quant_linear = QuantLinear(2, 4, bias=True, bias_quant=Int16Bias, return_quant_tensor=True)\n",
    "\n",
    "quant_output = quant_linear(float_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can solve the issue by passing in a `QuantTensor` coming from a different layer as input, or by setting an input quantizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[-0.6541,  0.1263,  0.1680, -0.1231],\n",
       "        [ 1.4658,  1.2395, -0.5207,  1.3989],\n",
       "        [ 1.6461,  0.8687, -1.0466,  1.4813]], grad_fn=<AddmmBackward>), scale=tensor([[9.0542e-05]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(18.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "float_input = torch.randn(3, 2)\n",
    "quant_linear = QuantLinear(\n",
    "    2, 4, bias=True, input_quant=Int8ActPerTensorFloat, bias_quant=Int16Bias, return_quant_tensor=True)\n",
    "\n",
    "quant_linear(float_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Operations on QuantTensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Element-wise Arithmetic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some basic arithmetic operations on QuantTensor are supported. Depending on the operation, some constraints can apply.  \n",
    "\n",
    "### Element-wise adds\n",
    "\n",
    "For element-wise addition, in line with traditional *fixed-point arithmetic*, the requirement is that scale of the operands is the same. \n",
    "\n",
    "In order to make it work with the way quantized activation operates by default during the statistics collection phase, this restriction is enforced in `eval()` mode but not in `train()` mode (during which the operands' scales are averaged).\n",
    "\n",
    "Calling the same QuantIdentity on tensors that need to be alligned is an easy way to do so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eval mode add quant inputs:\n",
      " QuantTensor(value=tensor([[ 1.5335, -0.2875],\n",
      "        [-2.0447,  0.5751],\n",
      "        [-1.0863, -1.4057]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(False)) \n",
      " QuantTensor(value=tensor([[ 0.3994,  0.8307],\n",
      "        [-0.7188, -0.3994],\n",
      "        [-0.5910,  0.1757]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(False)) \n",
      "\n",
      "Eval mode add quant output:\n",
      " QuantTensor(value=tensor([[ 1.9329,  0.5431],\n",
      "        [-2.7636,  0.1757],\n",
      "        [-1.6773, -1.2300]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(9.), signed_t=tensor(True), training_t=tensor(False))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "float_inp1 = torch.randn(3, 2)\n",
    "float_inp2 = torch.randn(3, 2)\n",
    "quant_identity = QuantIdentity(return_quant_tensor=True)\n",
    "\n",
    "#Training mode, statistics are being collected, scaling factors are different but it doesn't raise an error\n",
    "train_quant_inp1 = quant_identity(float_inp1)\n",
    "train_quant_inp2 = quant_identity(float_inp2)\n",
    "train_mode_add = train_quant_inp1 + train_quant_inp2\n",
    "\n",
    "#Inference mode, the EMA buffer is being used, scaling factors are the same\n",
    "quant_identity.eval()\n",
    "eval_quant_inp1 = quant_identity(float_inp1)\n",
    "eval_quant_inp2 = quant_identity(float_inp2)\n",
    "eval_mode_add = eval_quant_inp1 + eval_quant_inp2\n",
    "\n",
    "print(f\"Eval mode add quant inputs:\\n {eval_quant_inp1} \\n {eval_quant_inp2} \\n\")\n",
    "print(f\"Eval mode add quant output:\\n {eval_mode_add}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calling torch functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Through the `__torch_function__` interface (supported on PyTorch >= 1.5.0), standard torch functions can be called on a `QuantTensor`. Currently, for supported operations that are invariant to affine quantization a `QuantTensor` is returned, otherwise the output decays to a floating-point `torch.Tensor`. \n",
    "\n",
    "In this context for a function `func` to be invariant to quantization means that the output of the function applied to the dequantized value should still be a dequantized value with the same scale, zero-point, and bit-width, i.e. given an input dequantized value where `input_dequant_value / scale + zero_point == input_integer_value` that can be represented within `bit_width` bits, `func(input_dequant_value) / scale + zero_point == output_integer_value` that can also be represented within `bit_width` bits.\n",
    "\n",
    "### max_pool on QuantTensor\n",
    "\n",
    "For example, `torch.nn.functional.max_pool1d` is invariant to quantization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Quant input:\n",
      " QuantTensor(value=tensor([[[-1.1218, -1.1580, -0.2533, -0.4343],\n",
      "         [ 0.8504,  0.6876, -0.3076, -2.1170]],\n",
      "\n",
      "        [[ 0.4704, -0.1628,  1.4475,  0.2714],\n",
      "         [ 0.1628,  0.8685, -0.1448, -0.1086]],\n",
      "\n",
      "        [[ 0.9228,  1.2666,  2.0084,  0.0543],\n",
      "         [ 0.6152, -0.4162, -0.8323, -2.3160]]], grad_fn=<MulBackward0>), scale=tensor(0.0181, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True)) \n",
      "\n",
      "Quant output:\n",
      " QuantTensor(value=tensor([[[-1.1218, -0.2533],\n",
      "         [ 0.8504, -0.3076]],\n",
      "\n",
      "        [[ 0.4704,  1.4475],\n",
      "         [ 0.8685, -0.1086]],\n",
      "\n",
      "        [[ 1.2666,  2.0084],\n",
      "         [ 0.6152, -0.8323]]], grad_fn=<SqueezeBackward1>), scale=tensor(0.0181, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Alessandro\\miniconda3\\envs\\pt190\\lib\\site-packages\\torch\\nn\\functional.py:652: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  ..\\c10/core/TensorImpl.h:1156.)\n",
      "  return torch.max_pool1d(input, kernel_size, stride, padding, dilation, ceil_mode)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "float_inp = torch.randn(3, 2, 4)\n",
    "quant_identity = QuantIdentity(return_quant_tensor=True)\n",
    "\n",
    "quant_input = quant_identity(float_inp)\n",
    "quant_output = torch.nn.functional.max_pool1d(quant_input, kernel_size=2, stride=2)\n",
    "\n",
    "print(f\"Quant input:\\n {quant_input} \\n\")\n",
    "print(f\"Quant output:\\n {quant_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### tanh on QuantTensor\n",
    "\n",
    "While for example `torch.tanh` is not:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Quant input:\n",
      " QuantTensor(value=tensor([[[-1.1218, -1.1580, -0.2533, -0.4343],\n",
      "         [ 0.8504,  0.6876, -0.3076, -2.1170]],\n",
      "\n",
      "        [[ 0.4704, -0.1628,  1.4475,  0.2714],\n",
      "         [ 0.1628,  0.8685, -0.1448, -0.1086]],\n",
      "\n",
      "        [[ 0.9228,  1.2666,  2.0084,  0.0543],\n",
      "         [ 0.6152, -0.4162, -0.8323, -2.3160]]], grad_fn=<MulBackward0>), scale=tensor(0.0181, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True)) \n",
      "\n",
      "Quant output:\n",
      " tensor([[[-0.8082, -0.8204, -0.2480, -0.4089],\n",
      "         [ 0.6913,  0.5964, -0.2983, -0.9714]],\n",
      "\n",
      "        [[ 0.4386, -0.1614,  0.8952,  0.2649],\n",
      "         [ 0.1614,  0.7006, -0.1438, -0.1081]],\n",
      "\n",
      "        [[ 0.7272,  0.8529,  0.9646,  0.0542],\n",
      "         [ 0.5478, -0.3937, -0.6817, -0.9807]]], grad_fn=<TanhBackward>)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "float_inp = torch.randn(3, 2, 4)\n",
    "quant_identity = QuantIdentity(return_quant_tensor=True)\n",
    "\n",
    "quant_input = quant_identity(float_inp)\n",
    "quant_output = torch.tanh(quant_input)\n",
    "\n",
    "print(f\"Quant input:\\n {quant_input} \\n\")\n",
    "print(f\"Quant output:\\n {quant_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### QuantTensor concatenation\n",
    "\n",
    "Concatenation, i.e. `torch.cat`, can also be applied to quantized tensors whenever they all have the same sign, scale, zero-point and bit-width. Simiarly to element-wise adds, scale and zero-point are allowed to be different in training mode, but have to be the same in inference mode. Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eval mode concat quant inputs:\n",
      " QuantTensor(value=tensor([[ 1.5335, -0.2875],\n",
      "        [-2.0447,  0.5751],\n",
      "        [-1.0863, -1.4057]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(False)) QuantTensor(value=tensor([[ 0.3994,  0.8307],\n",
      "        [-0.7188, -0.3994],\n",
      "        [-0.5910,  0.1757]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(False)) \n",
      "\n",
      "Eval mode concat quant output:\n",
      " QuantTensor(value=tensor([[ 1.5335, -0.2875,  0.3994,  0.8307],\n",
      "        [-2.0447,  0.5751, -0.7188, -0.3994],\n",
      "        [-1.0863, -1.4057, -0.5910,  0.1757]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(False))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "float_inp1 = torch.randn(3, 2)\n",
    "float_inp2 = torch.randn(3, 2)\n",
    "quant_identity = QuantIdentity(return_quant_tensor=True)\n",
    "\n",
    "#Training mode, statistics are being collected, scaling factors are different but it doesn't raise an error\n",
    "train_mode_cat = torch.cat([quant_identity(float_inp1), quant_identity(float_inp2)], dim=1)\n",
    "\n",
    "#Inference mode, the EMA buffer is being used, scaling factors are the same\n",
    "quant_identity.eval()\n",
    "eval_quant_inp1 = quant_identity(float_inp1)\n",
    "eval_quant_inp2 = quant_identity(float_inp2)\n",
    "eval_mode_cat = torch.cat([eval_quant_inp1, eval_quant_inp2], dim=1)\n",
    "\n",
    "print(f\"Eval mode concat quant inputs:\\n {eval_quant_inp1} {eval_quant_inp2} \\n\")\n",
    "print(f\"Eval mode concat quant output:\\n {eval_mode_cat}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Customizing Quantizers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Common keyword arguments\n",
    "\n",
    "The easiest way to customize a quantizer is by passing appropriate keyword arguments. What happens any time we are setting a keyword argument is that we are overriding an hyperparameter in the underlying quantizer. Different quantizers expose different hyperparameters, which is why quantized layers like QuantLinear accept arbitrary `**kwargs`: \n",
    "\n",
    "### Weight bit-width\n",
    "\n",
    "Let's look at customizing the default quantizers first. To override the `bit_width` of weights in the default `weight_quant` quantizer we set `weight_bit_width`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weight QuantTensor:\n",
      " QuantTensor(value=tensor([[-0.0000,  0.3880],\n",
      "        [-0.5820, -0.5044],\n",
      "        [-0.2716,  0.1940],\n",
      "        [-0.0000,  0.5432]], grad_fn=<MulBackward0>), scale=tensor(0.0388, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(5.), signed_t=tensor(True), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "quant_linear = QuantLinear(2, 4, weight_bit_width=5, bias=True)\n",
    "\n",
    "print(f\"Weight QuantTensor:\\n {quant_linear.quant_weight()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Per-channel weight quantization\n",
    "\n",
    "We can enable channel-wise quantization by setting `weight_scaling_per_output_channel=True`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weight QuantTensor:\n",
      " QuantTensor(value=tensor([[-0.0000,  0.3793],\n",
      "        [-0.5820, -0.5044],\n",
      "        [-0.2723,  0.1816],\n",
      "        [-0.0000,  0.5607]], grad_fn=<MulBackward0>), scale=tensor([[0.0253],\n",
      "        [0.0388],\n",
      "        [0.0182],\n",
      "        [0.0374]], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(5.), signed_t=tensor(True), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "quant_linear = QuantLinear(2, 4, weight_bit_width=5, weight_scaling_per_output_channel=True, bias=False)\n",
    "\n",
    "print(f\"Weight QuantTensor:\\n {quant_linear.quant_weight()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Activation bit-width\n",
    "\n",
    "Similarly we can set the `bit_width` for activations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QuantTensor:\n",
      " QuantTensor(value=tensor([[ 1.6341, -0.5447],\n",
      "        [-2.1788,  0.5447],\n",
      "        [-1.0894, -1.6341]], grad_fn=<MulBackward0>), scale=tensor(0.5447, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(3.), signed_t=tensor(True), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "float_input = torch.randn(3, 2)\n",
    "quant_identity = QuantIdentity(bit_width=3, return_quant_tensor=True)\n",
    "print(f\"QuantTensor:\\n {quant_identity(float_input)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Activation quantization with max_val init\n",
    "\n",
    "Certain quantizers require the user to pass in additional keyword arguments. This is the case for example with `Uint8ActPerTensorFloatMaxInit`, where the learned scale factor is initialized by an user-defined maximum value rather than from statistics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[1.5294, 0.0000],\n",
       "        [0.0000, 0.5647],\n",
       "        [0.0000, 0.0000]], grad_fn=<MulBackward0>), scale=tensor(0.0235, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(False), training_t=tensor(True))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.quant import Uint8ActPerTensorFloatMaxInit\n",
    "\n",
    "float_inp1 = torch.randn(3, 2)\n",
    "quant_relu = QuantReLU(max_val=6.0, act_quant=Uint8ActPerTensorFloatMaxInit, return_quant_tensor=True)\n",
    "quant_relu(float_inp1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Per-channel activation quantization\n",
    "\n",
    "And its also possible to perform *per-channel activation quantization* if the number of channels is statically defined by providing some extra information. This makes sense for example as input to a depthwise separable layer, like a depthwise-separable convolution. Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Float input:\n",
      " tensor([[[ 1.5410, -0.2934, -2.1788,  0.5684, -1.0845],\n",
      "         [-1.3986,  0.4033,  0.8380, -0.7193, -0.4033]]]) \n",
      "\n",
      "Per-channel quant output:\n",
      " QuantTensor(value=tensor([[[ 0.8616, -0.7012,  0.4503],\n",
      "         [-1.1285, -0.4937, -0.1901]]], grad_fn=<SqueezeBackward1>), scale=tensor([[[0.0021],\n",
      "         [0.0013]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(17.), signed_t=tensor(True), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.nn import QuantConv1d\n",
    "\n",
    "BATCHES = 1\n",
    "CHANNELS = 2\n",
    "FEATURES = 5\n",
    "KERNEL = 3\n",
    "\n",
    "float_input = torch.randn(BATCHES, CHANNELS, FEATURES) \n",
    "per_channel_depthwise_quant_conv = QuantConv1d(\n",
    "    CHANNELS, CHANNELS, KERNEL, groups=CHANNELS, bias=True,\n",
    "    # set the quantizers\n",
    "    input_quant=Int8ActPerTensorFloat,\n",
    "    bias_quant=Int16Bias,\n",
    "    # customize the input quantizer with extra keyword args\n",
    "    input_bit_width=3,\n",
    "    input_scaling_per_output_channel=True, \n",
    "    input_scaling_stats_permute_dims=(1, 0, 2), # permute dims to put channels first during stats collection phase\n",
    "    input_per_channel_broadcastable_shape=(1, CHANNELS, 1), # shape of the learned parameter\n",
    "    return_quant_tensor=True)\n",
    "\n",
    "quant_output = per_channel_depthwise_quant_conv(float_input)\n",
    "\n",
    "print(f\"Float input:\\n {float_input} \\n\")\n",
    "print(f\"Per-channel quant output:\\n {quant_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inheriting from a quantizer\n",
    "\n",
    "What happens whenever we are passing a keyword argument (with prefix `weight_`, `input_`, `output_`, `bias_` for a layer like `QuantLinear`, no prefix for an activation layer like `QuantReLU`) is that we are overriding attributes of the underlying quantizer. To make things more compact and re-usable, we can also simply define a new quantizer by inheriting from the one we are customizing. The previous example then would look like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Float input:\n",
      " tensor([[[ 1.5410, -0.2934, -2.1788,  0.5684, -1.0845],\n",
      "         [-1.3986,  0.4033,  0.8380, -0.7193, -0.4033]]]) \n",
      "\n",
      "Per-channel quant output:\n",
      " QuantTensor(value=tensor([[[ 0.8616, -0.7012,  0.4503],\n",
      "         [-1.1285, -0.4937, -0.1901]]], grad_fn=<SqueezeBackward1>), scale=tensor([[[0.0021],\n",
      "         [0.0013]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(17.), signed_t=tensor(True), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.nn import QuantConv1d\n",
    "\n",
    "BATCHES = 1\n",
    "CHANNELS = 2\n",
    "FEATURES = 5\n",
    "KERNEL = 3\n",
    "\n",
    "class PerChannel3bActQuant(Int8ActPerTensorFloat):\n",
    "    bit_width = 3\n",
    "    scaling_per_output_channel=True\n",
    "    scaling_stats_permute_dims=(1, 0, 2)\n",
    "    \n",
    "float_input = torch.randn(BATCHES, CHANNELS, FEATURES) \n",
    "per_channel_depthwise_quant_conv = QuantConv1d(\n",
    "    CHANNELS, CHANNELS, KERNEL, groups=CHANNELS, bias=True,\n",
    "    # set the quantizers\n",
    "    input_quant=PerChannel3bActQuant,\n",
    "    bias_quant=Int16Bias,\n",
    "    # layer-specific kwarg\n",
    "    input_per_channel_broadcastable_shape=(1, CHANNELS, 1), \n",
    "    return_quant_tensor=True)\n",
    "\n",
    "quant_output = per_channel_depthwise_quant_conv(float_input)\n",
    "\n",
    "print(f\"Float input:\\n {float_input} \\n\")\n",
    "print(f\"Per-channel quant output:\\n {quant_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining a quantizer from scratch with enums "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far we have seen some options to customize the behaviour of existing quantizers. Let's look now at how exactly they are defined. \n",
    "\n",
    "Below is how two of the quantizers we have seen so far are declared within Brevitas using the enum-driven API (except for zero-point). As we will see in a second, it's not the only way to define a quantizer, but it's a way for users to easily experiment with different built-in options. \n",
    "\n",
    "The quantizer defines various directives and hyperparameters driving the underlying algorithms as class level attributes. \n",
    "\n",
    "## Weight quantizer \n",
    "\n",
    "For weights, What we are doing by inheriting from `WeightQuantSolver`, meaning those directive should be translated to an implementation that makes sense for weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brevitas.inject.enum import *\n",
    "from brevitas.core.zero_point import ZeroZeroPoint\n",
    "from brevitas.quant.solver import WeightQuantSolver, ActQuantSolver\n",
    "\n",
    "class Int8WeightPerTensorFloat(WeightQuantSolver):\n",
    "    quant_type = QuantType.INT # integer quantization\n",
    "    bit_width_impl_type = BitWidthImplType.CONST # constant bit width\n",
    "    float_to_int_impl_type = FloatToIntImplType.ROUND # round to nearest\n",
    "    scaling_impl_type = ScalingImplType.STATS # scale based on statistics\n",
    "    scaling_stats_op = StatsOp.MAX # scale statistics is the absmax value\n",
    "    restrict_scaling_type = RestrictValueType.FP # scale factor is a floating point value\n",
    "    scaling_per_output_channel = False # scale is per tensor\n",
    "    bit_width = 8 # bit width is 8\n",
    "    signed = True # quantization range is signed\n",
    "    narrow_range = True # quantization range is [-127,127] rather than [-128, 127]\n",
    "    zero_point_impl = ZeroZeroPoint # zero point is 0."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Activation quantizer \n",
    "\n",
    "Similarly for activations we can inherit from `ActQuantSolver`. The generated implementation is gonna be a bit different, since we are quantizing tensors defined at runtime rather than a parameter, by it's driven by the same enums:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Int8ActPerTensorFloat(ActQuantSolver):\n",
    "    quant_type = QuantType.INT # integer quantization\n",
    "    bit_width_impl_type = BitWidthImplType.CONST # constant bit width\n",
    "    float_to_int_impl_type = FloatToIntImplType.ROUND # round to nearest\n",
    "    scaling_impl_type = ScalingImplType.PARAMETER_FROM_STATS # scale is a parameter initialized from statistics\n",
    "    scaling_stats_op = StatsOp.PERCENTILE # scale statistics is a percentile of the abs value \n",
    "    high_percentile_q = 99.999 # percentile is 99.999\n",
    "    collect_stats_steps = 300  # statistics are collected for 300 forward steps before switching to a learned parameter\n",
    "    restrict_scaling_type = RestrictValueType.FP # scale is a floating-point value\n",
    "    scaling_per_output_channel = False  # scale is per tensor\n",
    "    bit_width = 8  # bit width is 8\n",
    "    signed = True # quantization range is signed\n",
    "    narrow_range = False # quantization range is [-128, 127] rather than [-127, 127]\n",
    "    zero_point_impl = ZeroZeroPoint # zero point is 0."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Any of attributes in the quantizers above can be passed in or overriden as a keyword argument to the layer the quantizer is passed to (together with its appropriate prefix), with keyword arguments always having the priority over the value defined in the quantizer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learned scale and bit-width quantizer\n",
    "\n",
    "Let's now look at how we can build something more advanced by tweaking a few options. We want a weight quantizer with: \n",
    "- per-channel scale factors learned in log domain as a parameter initialized from absmax statistics.\n",
    "- bit-width initialized to 8b and learned as a parameter from there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weight QuantTensor:\n",
      " QuantTensor(value=tensor([[-0.0060,  0.3793],\n",
      "        [-0.5820, -0.5224],\n",
      "        [-0.2723,  0.1887],\n",
      "        [-0.0132,  0.5607]], grad_fn=<MulBackward0>), scale=tensor([[0.0030],\n",
      "        [0.0046],\n",
      "        [0.0021],\n",
      "        [0.0044]], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8., grad_fn=<RoundSteFnBackward>), signed_t=tensor(True), training_t=tensor(True))\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.quant import Int8WeightPerTensorFloat\n",
    "\n",
    "class LearnedIntWeightPerChannelFloat(Int8WeightPerTensorFloat):\n",
    "    scaling_per_output_channel = True\n",
    "    scaling_impl_type = ScalingImplType.PARAMETER_FROM_STATS\n",
    "    restrict_scaling_type = RestrictValueType.LOG_FP\n",
    "    bit_width_impl_type = BitWidthImplType.PARAMETER \n",
    "\n",
    "    \n",
    "quant_linear = QuantLinear(2, 4, weight_quant=LearnedIntWeightPerChannelFloat, bias=False)\n",
    "\n",
    "print(f\"Weight QuantTensor:\\n {quant_linear.quant_weight()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the quantized weights' `bit_width` is now a value we can backpropagate through, as it exposes a `grad_fn` function. This way we can use it as part of a loss regularization function to model a particular hardware cost function, e.g. pushing larger layers to have a smaller bit width. The same principle applies to activation, and we can combine them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[-0.9109, -0.4588,  0.3119, -0.6530],\n",
       "        [ 1.2089,  0.6493, -0.3731,  0.8706],\n",
       "        [ 1.3893,  0.2823, -0.8979,  0.9543]], grad_fn=<MmBackward>), scale=tensor([[9.0542e-05, 3.9068e-05, 5.6866e-05, 6.4251e-05]],\n",
       "       grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(17., grad_fn=<CeilSteFnBackward>), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "class LearnedIntActPerTensorFloat(Int8ActPerTensorFloat):\n",
    "    bit_width_impl_type = BitWidthImplType.PARAMETER \n",
    "    restrict_scaling_type = RestrictValueType.LOG_FP\n",
    "\n",
    "float_inp = torch.randn(3, 2)\n",
    "quant_linear = QuantLinear(\n",
    "    2, 4, \n",
    "    input_quant=LearnedIntActPerTensorFloat,\n",
    "    weight_quant=LearnedIntWeightPerChannelFloat, \n",
    "    return_quant_tensor=True, bias=False)\n",
    "\n",
    "quant_linear(float_inp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, we can backpropagate through the output `bit_width`. By including it in a loss function, we could then try to control the size of the output accumulator, and in turn the `bit_width` of both weights and input."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retraining from floating-point"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In many scenarios it's convenient to perform quantization-aware training starting from a floating-point model. Say for example we want to load a pretrained-floating point state on top of the layer with learned bit-width and scale we just saw. We simulate it with a separate floating-point `nn.Linear`. If we didn't do anything else we would get an error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Error(s) in loading state_dict for QuantLinear:\n\tMissing key(s) in state_dict: \"input_quant.fused_activation_quant_proxy.tensor_quant.scaling_impl.value\", \"input_quant.fused_activation_quant_proxy.tensor_quant.msb_clamp_bit_width_impl.bit_width_offset\", \"weight_quant.tensor_quant.scaling_impl.value\", \"weight_quant.tensor_quant.msb_clamp_bit_width_impl.bit_width_offset\". ",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32mC:\\Users\\ALESSA~1\\AppData\\Local\\Temp/ipykernel_18920/1653109852.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     10\u001b[0m     return_quant_tensor=True, bias=False)\n\u001b[0;32m     11\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 12\u001b[1;33m \u001b[0mquant_linear\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfloat_linear\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mstate_dict\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32m~\\miniconda3\\envs\\pt190\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36mload_state_dict\u001b[1;34m(self, state_dict, strict)\u001b[0m\n\u001b[0;32m   1405\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0merror_msgs\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m>\u001b[0m \u001b[1;36m0\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1406\u001b[0m             raise RuntimeError('Error(s) in loading state_dict for {}:\\n\\t{}'.format(\n\u001b[1;32m-> 1407\u001b[1;33m                                self.__class__.__name__, \"\\n\\t\".join(error_msgs)))\n\u001b[0m\u001b[0;32m   1408\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0m_IncompatibleKeys\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmissing_keys\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0munexpected_keys\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1409\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: Error(s) in loading state_dict for QuantLinear:\n\tMissing key(s) in state_dict: \"input_quant.fused_activation_quant_proxy.tensor_quant.scaling_impl.value\", \"input_quant.fused_activation_quant_proxy.tensor_quant.msb_clamp_bit_width_impl.bit_width_offset\", \"weight_quant.tensor_quant.scaling_impl.value\", \"weight_quant.tensor_quant.msb_clamp_bit_width_impl.bit_width_offset\". "
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from torch import nn\n",
    "\n",
    "float_linear = nn.Linear(2, 4, bias=False)\n",
    "quant_linear = QuantLinear(\n",
    "    2, 4, \n",
    "    input_quant=LearnedIntActPerTensorFloat,\n",
    "    weight_quant=LearnedIntWeightPerChannelFloat, \n",
    "    return_quant_tensor=True, bias=False)\n",
    "\n",
    "quant_linear.load_state_dict(float_linear.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's because the quantizers for both weights and input are introducing new learned parameters that are not present in the original floating-point layer.\n",
    "\n",
    "We can supress the error to accomodate the re-training scenario by either setting the env variable `BREVITAS_IGNORE_MISSING_KEYS=1`, or by enabling the corresponding config flag:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from torch import nn\n",
    "from brevitas import config\n",
    "\n",
    "config.IGNORE_MISSING_KEYS = True\n",
    "\n",
    "float_linear = nn.Linear(2, 4, bias=False)\n",
    "quant_linear = QuantLinear(\n",
    "    2, 4, \n",
    "    input_quant=LearnedIntActPerTensorFloat,\n",
    "    weight_quant=LearnedIntWeightPerChannelFloat, \n",
    "    return_quant_tensor=True, bias=False)\n",
    "\n",
    "quant_linear.load_state_dict(float_linear.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, in this scenario we are saying we want the learned component of the scale factor for weights to be initialized to the maximum value of the weight tensor on a per channel basis. So even though it's not part of the floating-point state dict, the value of the weight scale factor should be updated. Thanks to how the weight quantizer works this happens *automatically*:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weight scale before load_state_dict: tensor([[0.0015],\n",
      "        [0.0017],\n",
      "        [0.0053],\n",
      "        [0.0023]], grad_fn=<DivBackward0>)\n",
      "Weight scale after load_state_dict: tensor([[0.0030],\n",
      "        [0.0046],\n",
      "        [0.0021],\n",
      "        [0.0044]], grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from torch import nn\n",
    "from brevitas import config\n",
    "\n",
    "config.IGNORE_MISSING_KEYS = True\n",
    "\n",
    "float_linear = nn.Linear(2, 4, bias=False)\n",
    "quant_linear = QuantLinear(\n",
    "    2, 4, \n",
    "    input_quant=LearnedIntActPerTensorFloat,\n",
    "    weight_quant=LearnedIntWeightPerChannelFloat, \n",
    "    return_quant_tensor=True, bias=False)\n",
    "\n",
    "print(f'Weight scale before load_state_dict: {quant_linear.quant_weight().scale}')\n",
    "\n",
    "quant_linear.load_state_dict(float_linear.state_dict())\n",
    "\n",
    "print(f'Weight scale after load_state_dict: {quant_linear.quant_weight().scale}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This works also in the scenario where the weight quantizer is shared among multiple layers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining a quantizer from scratch with dependency-injection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far we have seen how users can pick and adopt various quantization algorithms that can be expressed through components already built into Brevitas. The more interesting question is how to expressed completely novel algorithms based on new quantization components.\n",
    "\n",
    "## Activation quantization from scratch\n",
    "\n",
    "To do so we need to move away from the enum-driven API, and understand how quantizers work underneath. If we look at the structure of a layer like `QuantReLU`, we can immediately notice that it's formed by a variety of nested modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantReLU(\n",
       "  (input_quant): ActQuantProxyFromInjector(\n",
       "    (_zero_hw_sentinel): StatelessBuffer()\n",
       "  )\n",
       "  (act_quant): ActQuantProxyFromInjector(\n",
       "    (_zero_hw_sentinel): StatelessBuffer()\n",
       "    (fused_activation_quant_proxy): FusedActivationQuantProxy(\n",
       "      (activation_impl): ReLU()\n",
       "      (tensor_quant): RescalingIntQuant(\n",
       "        (int_quant): IntQuant(\n",
       "          (float_to_int_impl): RoundSte()\n",
       "          (tensor_clamp_impl): TensorClamp()\n",
       "          (delay_wrapper): DelayWrapper(\n",
       "            (delay_impl): _NoDelay()\n",
       "          )\n",
       "        )\n",
       "        (scaling_impl): ParameterFromRuntimeStatsScaling(\n",
       "          (stats_input_view_shape_impl): OverTensorView()\n",
       "          (stats): _Stats(\n",
       "            (stats_impl): AbsPercentile()\n",
       "          )\n",
       "          (restrict_clamp_scaling): _RestrictClampValue(\n",
       "            (clamp_min_ste): Identity()\n",
       "            (restrict_value_impl): FloatRestrictValue()\n",
       "          )\n",
       "          (restrict_inplace_preprocess): Identity()\n",
       "          (restrict_preprocess): Identity()\n",
       "        )\n",
       "        (int_scaling_impl): IntScaling()\n",
       "        (zero_point_impl): ZeroZeroPoint(\n",
       "          (zero_point): StatelessBuffer()\n",
       "        )\n",
       "        (msb_clamp_bit_width_impl): BitWidthConst(\n",
       "          (bit_width): StatelessBuffer()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "QuantReLU()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Where the implementation of the quantization component is expressed by the `tensor_quant` module.\n",
    "\n",
    "The idea is that we express a quantization algorithm as various `nn.Module` combined in a depedency-injection fashion, meaning that they combine from the outside-in through standardized interfaces. \n",
    "\n",
    "So we have different modules to express a particular type of scale factor, a particular type of rounding, a particular integer quantization procedure, and so on. The overall quantization algorithm is then expressed as a nested assembly of all its different pieces. If we were to do it manually for the layer above, it would look like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brevitas.core.function_wrapper import TensorClamp, OverTensorView, RoundSte\n",
    "from brevitas.core.scaling import ParameterFromRuntimeStatsScaling, IntScaling\n",
    "from brevitas.core.stats import AbsPercentile\n",
    "from brevitas.core.restrict_val import FloatRestrictValue\n",
    "from brevitas.core.bit_width import BitWidthConst\n",
    "from brevitas.core.quant import IntQuant, RescalingIntQuant\n",
    "from brevitas.core.zero_point import ZeroZeroPoint\n",
    "\n",
    "tensor_quant = RescalingIntQuant(\n",
    "    int_quant=IntQuant(\n",
    "        float_to_int_impl=RoundSte(),\n",
    "        tensor_clamp_impl=TensorClamp(),\n",
    "        signed=False,\n",
    "        narrow_range=False),\n",
    "    zero_point_impl=ZeroZeroPoint(),\n",
    "    bit_width_impl=BitWidthConst(bit_width=8),\n",
    "    int_scaling_impl=IntScaling(signed=False, narrow_range=False),\n",
    "    scaling_impl=ParameterFromRuntimeStatsScaling(\n",
    "        scaling_stats_input_view_shape_impl=OverTensorView(),\n",
    "        scaling_stats_impl=AbsPercentile(\n",
    "            high_percentile_q=99.999, stats_reduce_dim=None),\n",
    "        collect_stats_steps=300)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this approach we can express arbitrary quantization algorithms in a very modular fashion. \n",
    "However, we have two extra constraints we care about:\n",
    "\n",
    "- This degree of modularity comes at a cost, i.e. the overhead of doing a lot of small operations scattered across different modules in Python shows. We want to avoid it by taking advantage of PyTorch's TorchScript JIT compiler, meaning the `tensor_quant` module above should be fused and end-to-end compiled.\n",
    "\n",
    "- We want a way to define our quantizer such that it can be re-instantiated and re-compiled on-demand in certain scenarios, like whenever it's shared among multiple layers or when a pretrained `state_dict` is load on top of it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To solve the first issue, we restrict ourselves to using only `ScriptModule`, a variant of a `nn.Module` that can be JIT-compiled. This is also why we favour composition over inheritance, as TorchScript has almost no support for inheritance. So all the modules we are instantiating above derive from an alias to a ScriptModule called `brevitas.jit.ScriptModule`, whose compilation can be enabled by setting the env variable `BREVITAS_JIT=1`.\n",
    "\n",
    "To solve the second issue, we adopt and extend an auto-wiring dependency injection library called *[dependencies](https://github.com/proofit404/dependencies)* that performs the assembly and instantiation process of `tensor_quant` automatically. Let's look at an example to make things more clear:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brevitas.inject import ExtendedInjector\n",
    "from brevitas.proxy import ActQuantProxyFromInjector\n",
    "\n",
    "class Int8ActPerTensorFloatFromScratch(ExtendedInjector):\n",
    "    proxy_class = ActQuantProxyFromInjector # Python wrapper, returns QuantTensor\n",
    "    tensor_quant = RescalingIntQuant # TorchScript implementation, returns (dequant_value, scale, zero_point, bit_width)\n",
    "    int_quant = IntQuant\n",
    "    float_to_int_impl = RoundSte \n",
    "    zero_point_impl = ZeroZeroPoint\n",
    "    bit_width_impl = BitWidthConst\n",
    "    tensor_clamp_impl = TensorClamp\n",
    "    int_scaling_impl = IntScaling\n",
    "    scaling_impl = ParameterFromRuntimeStatsScaling\n",
    "    scaling_stats_input_view_shape_impl = OverTensorView\n",
    "    scaling_stats_impl = AbsPercentile\n",
    "    restrict_scaling_impl = FloatRestrictValue\n",
    "    high_percentile_q = 99.999 \n",
    "    collect_stats_steps = 300\n",
    "    scaling_shape = ()\n",
    "    stats_reduce_dim = None\n",
    "    bit_width = 8\n",
    "    narrow_range = True\n",
    "    signed = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Int8ActPerTensorFloatFromScratch` is equivalent to `Int8ActPerTensorFloat`, but defined directly in terms of the modules that define the implementation, rather than with enums. As such, we inherit from our dependency-injection base class, the `ExtendedInjector`. The quantizer needs to define two attributes: \n",
    "- `proxy_class` represents the Python wrapper that interfaces a quantized layer like QuantIdentity with the quantizer itself, and for convinience is defined as part of the quantizer itself. It has to return a `QuantTensor`. \n",
    "- `tensor_quant` is the actual ScriptModule implementating our quantization algorithm. It has to return a 4-tuple containing tensors (dequant_value, scale, zero_point, bit_width).\n",
    "\n",
    "`tensor_quant` is automatically instantiated by the dependency-injection engine by matching its keyword argument to attributes of the `ÈxtendedInjector`. If one of the attributes is a class that requires to be instantiated the process is repeated in a recursive manner. If it's already an instance, it's taken as is.\n",
    "\n",
    "Note that the enum-driven API only adds an extra layer of abstraction on top of this, but it maps to the same components with the same names.\n",
    "\n",
    "This way, indipendently of how it's originally defined, a user can override any component of an existing quantizer in a fine grained manner, or define a new quantizer from scratch with entirely custom components. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Weight quantization with learned scale from scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, we can define a weight quantizer from scratch with a learned scale factor implemented by `ParameterScaling` module. To initialize it, we define `scaling_init` as a function of the full precision weights to quantize, by taking the absmax of weight tensor, as with did in previous examples. To do so, we define a `@value` function, a function which is evaluated at *dependency-injection time*:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brevitas.inject import value\n",
    "from brevitas.proxy import WeightQuantProxyFromInjector\n",
    "from brevitas.core.scaling import ParameterScaling\n",
    "\n",
    "class Int8ActPerTensorFloatParameterFromScratch(ExtendedInjector):\n",
    "    \n",
    "    @value\n",
    "    def scaling_init(module):\n",
    "        return module.weight.abs().max().detach()\n",
    "    \n",
    "    proxy_class = WeightQuantProxyFromInjector # Python wrapper, returns QuantTensor\n",
    "    tensor_quant = RescalingIntQuant \n",
    "    int_quant = IntQuant\n",
    "    float_to_int_impl = RoundSte \n",
    "    zero_point_impl = ZeroZeroPoint\n",
    "    bit_width_impl = BitWidthConst\n",
    "    tensor_clamp_impl = TensorClamp\n",
    "    int_scaling_impl = IntScaling\n",
    "    scaling_impl = ParameterScaling\n",
    "    restrict_scaling_impl = FloatRestrictValue\n",
    "    scaling_shape = ()\n",
    "    bit_width = 8\n",
    "    narrow_range = True\n",
    "    signed = True\n",
    "    \n",
    "quant_linear = QuantLinear(2, 4, weight_quant=Int8ActPerTensorFloatParameterFromScratch, bias=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is how `scaling_impl_type = ScalingImplType.PARAMETER_FROM_STATS` for weight quantization works underneath. \n",
    "\n",
    "the argument `module` of `scaling_init` represents the layer (or tuple of layers) the quantizer is passed to, and it allows the quantizer to depend on properties of the layer to quantize.\n",
    "\n",
    "By declaring `scaling_init` in this way we can easily recompute it any time the layer changes. As we saw previously, for changes to the `state_dict` this happens automatically. For in-place changes like weight initialization, which cannot easily be intercepted automatically, we can invoke it manually by calling `quant_linear.weight_quant.init_tensor_quant()`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sharing learned bit-width among layers\n",
    "\n",
    "We saw in a previous section how to share entire quantizers among layers. We are now interested in sharing only certain components, and we leverage the dependency-injection mechanism we just saw to do so. \n",
    "\n",
    "For example, we want the `QuantIdentity` and the `QuantLinear` with learned bit-width of our previous example to share the same learned bit-width. That can be useful whenever a particular hardware target supports only operations among tensors of the same datatype . With some verbosity, doing so is possible:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "float_inp = torch.randn(3, 2)\n",
    "quant_linear = QuantLinear(\n",
    "    2, 4, weight_quant=LearnedIntWeightPerChannelFloat, return_quant_tensor=True, bias=False)\n",
    "quant_identity = QuantIdentity(\n",
    "    bit_width_impl=quant_linear.weight_quant.tensor_quant.msb_clamp_bit_width_impl,\n",
    "    act_quant=LearnedIntActPerTensorFloat, return_quant_tensor=True)\n",
    "\n",
    "quant_identity_bit_width = quant_identity.act_quant.fused_activation_quant_proxy.tensor_quant.msb_clamp_bit_width_impl\n",
    "quant_linear_bit_width = quant_linear.weight_quant.tensor_quant.msb_clamp_bit_width_impl\n",
    "\n",
    "quant_identity_bit_width is quant_linear_bit_width"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see the two bit-width implementations are the same instance. What we are doing is sharing the instance of `msb_clamp_bit_width_impl` from one quantizer to another. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Export"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Brevitas doesn't perform any kind of acceleration on its own. To do so, a quantized model needs to be exported to a backend first. In a perfect world there would be a single way to represent quantized models. In practice however that's not the case, and different toolchains make different assumptions on what kind of quantization is supported and how it should be represented. \n",
    "\n",
    "Brevitas then supports different flavours of export at different levels of abstractions by taking advantage of PyTorch's support for custom symbolic representations, and in particular ONNX.\n",
    "\n",
    "Existing export flows assume static quantization, i.e. scales, zero-point and bit-width need to be indipedentent from the input. All export flows abstract away from the specifics of how scales, zero-point and bit-width have been determined. \n",
    "However, different export flows provide supports only for certain combinations of scales, zero-point, precision, or structure of the model.\n",
    "\n",
    "We install onnx and onnxoptimizer as they are required for export, and we use Netron to visualize the exported models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: netron in c:\\users\\alessandro\\miniconda3\\envs\\pt190\\lib\\site-packages (5.3.9)\n",
      "Requirement already satisfied: onnx in c:\\users\\alessandro\\miniconda3\\envs\\pt190\\lib\\site-packages (1.10.2)\n",
      "Requirement already satisfied: onnxoptimizer in c:\\users\\alessandro\\miniconda3\\envs\\pt190\\lib\\site-packages (0.2.6)\n",
      "Requirement already satisfied: numpy>=1.16.6 in c:\\users\\alessandro\\miniconda3\\envs\\pt190\\lib\\site-packages (from onnx) (1.21.2)\n",
      "Requirement already satisfied: typing-extensions>=3.6.2.1 in c:\\users\\alessandro\\miniconda3\\envs\\pt190\\lib\\site-packages (from onnx) (3.10.0.2)\n",
      "Requirement already satisfied: protobuf in c:\\users\\alessandro\\miniconda3\\envs\\pt190\\lib\\site-packages (from onnx) (3.19.1)\n",
      "Requirement already satisfied: six in c:\\users\\alessandro\\miniconda3\\envs\\pt190\\lib\\site-packages (from onnx) (1.16.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install netron onnx onnxoptimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "import netron\n",
    "import time\n",
    "from IPython.display import IFrame\n",
    "\n",
    "def show_netron(model_path, port):\n",
    "    time.sleep(3.)\n",
    "    netron.start(model_path, address=(\"localhost\", port), browse=False)\n",
    "    return IFrame(src=f\"http://localhost:{port}/\", width=\"100%\", height=400)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export to ONNX QOps\n",
    "\n",
    "Say we want to export a QuantConv1d with 4b symmetric weights, 8b symmetric inputs and outputs, and 16 biases. \n",
    "We can export it to a ONNX's `QLinearConv`, but some information will be lost. In particular, weights will be represented as 8b and bias as 32b, even though they are respectively 4b and 16b. This is because ONNX does not provide a standardized way to represent them as such:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.nn import QuantConv1d\n",
    "from brevitas.quant import Int8WeightPerTensorFloat, Int8ActPerTensorFloat, Int16Bias\n",
    "from brevitas.export import export_onnx_qop\n",
    "\n",
    "float_inp = torch.randn(1, 2, 5)\n",
    "\n",
    "quant_conv_4b8b = QuantConv1d(\n",
    "    2, 4, 3, bias=True, weight_bit_width=4,\n",
    "    input_quant=Int8ActPerTensorFloat,\n",
    "    output_quant=Int8ActPerTensorFloat,\n",
    "    bias_quant=Int16Bias)\n",
    "\n",
    "output_path = 'qop_onnx_conv_4b8b.onnx'\n",
    "export_onnx_qop(quant_conv_4b8b, input_t=float_inp, export_path=output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'qop_onnx_conv_4b8b.onnx' at http://localhost:8082\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8082/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x1720d689b38>"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(output_path, 8082)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general the standard ONNX opset doesn't support representing quantization below 8b. Additionally, ONNX QOp representation requires an output quantizer to be set at part of of the layer. \n",
    "\n",
    "The constraint of always having an output quantizer is relaxed in the more recently introduced QDQ style of representation (for which there is support in Brevitas starting from version 0.8), which uses only `QuantizeLinear` and `DequantizeLinear` to represent quantization, but even with that support is still limited to 8b quantization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export to custom Quantized ONNX\n",
    "\n",
    "As an alternative, we can export it to QONNX, a custom ONNX dialect that Brevitas defines with support for custom quantization operators that can capture those informations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "from brevitas.export import export_qonnx\n",
    "from brevitas.quant import Int8WeightPerTensorFloat, Int8ActPerTensorFloat, Int16Bias\n",
    "\n",
    "float_inp = torch.randn(1, 2, 5)\n",
    "\n",
    "quant_conv_4b8b = QuantConv1d(\n",
    "    2, 4, 3, bias=True, weight_bit_width=4,\n",
    "    input_quant=Int8ActPerTensorFloat,\n",
    "    output_quant=Int8ActPerTensorFloat,\n",
    "    bias_quant=Int16Bias)\n",
    "\n",
    "output_path = 'brevitas_onnx_conv4b8b.onnx'\n",
    "export_qonnx(quant_conv_4b8b, input_t=float_inp, export_path=output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'brevitas_onnx_conv4b8b.onnx' at http://localhost:8083\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8083/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x1720e7cf128>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(output_path, 8083)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the `Quant` nodes above, arbitrary scale, zero-point and bit-widths are supported. This way we can support exporting scenarios where, for example, we are quantizing only weights to 4b, which a QLinearConv wouldn't be able to capture:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "\n",
    "quant_conv_4b_weights = QuantConv1d(2, 4, 3, bias=True, weight_bit_width=4)\n",
    "\n",
    "output_path = 'brevitas_onnx_conv_4b_weights.onnx'\n",
    "export_qonnx(quant_conv_4b_weights, input_t=float_inp, export_path=output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'brevitas_onnx_conv_4b_weights.onnx' at http://localhost:8084\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8084/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x1720e80a978>"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(output_path, 8084)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The custom format shown above can integrated into ONNX-based toolchains, e.g. it's supported by our own FINN toolchain for low-precision dataflow style custom FPGAs implementations, and would be a starting point for direct integration with TVM.\n",
    "\n",
    "## Export to TorchScript quantization backend\n",
    "\n",
    "It's also possible to export to TorchScript own quantized functional operators, which come with their own set of restrictions. In particular, weights should be 7b and unsigned, which requires a zero-point. We can model that with appropriate quantizers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brevitas.quant import ShiftedUint8ActPerTensorFloat\n",
    "from brevitas.export import export_torch_qop\n",
    "\n",
    "\n",
    "quant_conv_8b7b = QuantConv1d(\n",
    "    2, 4, 3, bias=True,\n",
    "    input_quant=ShiftedUint8ActPerTensorFloat,\n",
    "    output_quant=ShiftedUint8ActPerTensorFloat,\n",
    "    weight_bit_width=7,\n",
    "    bias_quant=Int16Bias)\n",
    "\n",
    "output_path = 'pytorch_qf_conv_8b7b.pt'\n",
    "export_torch_qop(quant_conv_8b7b, input_t=float_inp, export_path=output_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\users\\alessandro\\documenti\\brevitas_tvmcon\\src\\brevitas\\quant_tensor\\__init__.py:74: TracerWarning: torch.tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.\n",
      "  training = torch.tensor(training, dtype=torch.bool)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Serving 'pytorch_qf_conv_8b7b.pt' at http://localhost:8085\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"400\"\n",
       "            src=\"http://localhost:8085/\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x1720e87a438>"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_netron(output_path, 8085)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see though information on the fact that activations are 7b is lost, and they simply marked as 8b.\n",
    "\n",
    "Additionally, because bias quantization is not represented explicitly (although it is performed implicitly at 32b at runtime in the backend), any information around that is lost.\n",
    "As with standard ONNX, representing precisions below 8b is not possible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Brevitas and FX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we have covered how to incorporate quantized layers into the design of a neural networks. However, in some scenarios it might convenient to start from an existing definition of a neural networks in floating-point, and programmatically replace floating-point layers with quantized layers according to some patterns.\n",
    "\n",
    "PyTorch's recently introduced FX graph subsystem enables this kind of transformations. While covering FX goes beyond the scope of this tutorial, it's worth mentioning that Brevitas has embraced FX, and it actually implements a backport of PyTorch's 1.8 FX (the current LTS) under `brevitas.fx`, together with a custom input-driven tracer (`brevitas.fx.value_trace`) which similarly to `torch.jit.trace` allows to trace through conditionals and unpacking operations, as well as various graph transformations under `brevitas.graph`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calibration-based post-training quantization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While Brevitas was originally designed for quantization-aware training, performing calibration-based post-training quantization is possible, and it can be useful on its own or an intermediate step between floating-point training and quantization-aware retraining. \n",
    "\n",
    "The idea of calibration-based quantization is to performs forward passes only with a small set of data and collect statistics to determine scale factors and zero-points. \n",
    "\n",
    "Any Brevitas quantizer that is based on statistics can be used for this purpose, with the caveat that don't want quantization to be enabled while statistics are being collected, or the data wouldn't be representative of what the floating-point model does, so we temporarely disabling quantization while doing so. Afterwards, we re-enable calibration and apply bias correction. \n",
    "\n",
    "Assuming `quant_model` is a quantized model, with pretrained floating-point weights loaded on top, we can do as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "from brevitas.graph.calibrate import bias_correction_mode\n",
    "from brevitas.graph.calibrate import calibration_mode\n",
    "\n",
    "def calibrate_model(calibration_loader, quant_model, args):\n",
    "    with torch.no_grad():\n",
    "        # Put the model in calibration mode to collect statistics\n",
    "        # Quantization is automatically disabled during the calibration, and re-enabled at the end\n",
    "        with calibration_mode(quant_model):\n",
    "            for i, (images, _) in enumerate(calibration_loader):\n",
    "                print(f'Calibration iteration {i}')\n",
    "                quant_model(images)\n",
    "        \n",
    "        # Apply bias correction\n",
    "        with bias_correction_mode(quant_model):\n",
    "            for i, (images, _) in enumerate(calibration_loader):\n",
    "                print(f'Bias correction iteration {i}')\n",
    "                quant_model(images)\n",
    "    return quant_model"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch_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.8.15 | packaged by conda-forge | (default, Nov 22 2022, 08:42:03) [MSC v.1929 64 bit (AMD64)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "b6e150ee02c45d2c3f896173a651a21b25567e05411969bcc0f3a62fa15a0a0b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}