tutorials/quantization/basic.ipynb (437 lines of code) (raw):
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Quantization\n",
"\n",
"You may all know that quantization leads to smaller model size and faster model inference. But do you know why? Here we will cover the basics of quantization.\n",
"\n",
"## Basic concept\n",
"For a floating value `f`, it can be expressed as an integral number `q = f / s + o` (aka quantized value) given the quantization parameters `s` as scale and `o` as offset.\n",
"When you convert `f` to `q`, the action is called \"Quantize\". Respectively, if you restore `f` from `q`, it is called \"Dequantize\".\n",
"\n",
"Usually, we use 8-bit quantization, which means the value of `q` is within range `[0, 255]` (unsigned) or `[-128, 127]` (signed).\n",
"\n",
"## Dynamic quantization & static quantization\n",
"Generally speaking, there are two kinds of quantization. Dynamic quantization (aka hybrid quantization or dynamic range quantization) is a kind of quantization that calculates the quantization parameters on the fly. All you need to do is to convert the weights to quantized values. You may refer to [this tutorial](../hybrid.ipynb) for performing this kind of quantization using TinyNerualNetwork.\n",
"\n",
"Static quantization, on the other hand, have the quantization parameters calculated before the inference phase. There are generally two ways to achieve that, [quantization aware training](../qat.ipynb) and [post quantization](../post.ipynb). We will illustrate the process of it with the example in the next section.\n",
"\n",
"If you want to choose one type of quantization without knowing the details, you may base on the decision tree in the graph or the summary table below.\n",
"\n",
"\n",
"\n",
"| Technique | Benefits | Hardware |\n",
"|----------------------------|------------------------------|---------------------------------|\n",
"| Dynamic quantization | 4x smaller, 2x-3x speedup | CPU |\n",
"| Static quantization | 4x smaller, 3x+ speedup | CPU, Edge TPU, Microcontrollers |\n",
"\n",
"## How static quantization is performed in DNN frameworks?\n",
"The key here is fake quantization. What is fake quantization? Suppose we have only one operation `y = conv(x)` in the original floating computation graph, then we want to have `y’ = q_conv(x‘)` in the quantized graph. \n",
"\n",
"With fake quantization, we have `x’ = fake_quantize(x)` and `y’ = fake_quantize(y)`.\n",
"First, we will observe the mininum and maximum values of `x` and `y`. Let's mark them as `x_min`, `x_max` and `y_min`, `y_max`. And then we calculate the quantization parameters, including scale `s` and offset `o`.\n",
"\n",
"Asymmetric quantization:\n",
"```py\n",
"s = (f_max - f_min) / (q_max - q_min)\n",
"o = q_min - min(f_min, 0) / s\n",
"```\n",
"\n",
"Symmetric quantization:\n",
"```py\n",
"s = max(f_max, -fmin) / ((q_max - q_min) / 2)\n",
"o = 128 [uint8]\n",
"o = 0 [int8]\n",
"```\n",
"\n",
"Then, fake quantization is performed using the given quantization parameters. We have `x’ = fake_quantize(x) = (clamp(round(x / s + o), q_min, q_max) - o) * s`. Similarly, we can get `y‘`.\n",
"\n",
"Finally, we replace the floating kernels with the quantized kernels. So the computation graph will contain the following operations.\n",
"```py\n",
"x’ = quantize(x, s_x, o_x)\n",
"y’ = q_conv(x)\n",
"y = dequantize(y’, s_y, o_y)\n",
"```\n",
"\n",
"## Static quantization in PyTorch\n",
"We use the following PyTorch model as an example."
]
},
{
"cell_type": "code",
"execution_count": 28,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
"\n",
"class Model(nn.Module):\n",
" def __init__(self):\n",
" super().__init__()\n",
" self.conv = nn.Conv2d(3, 3, 1)\n",
" self.bn = nn.BatchNorm2d(3)\n",
" self.relu = nn.ReLU()\n",
" \n",
" def forward(self, x):\n",
" y = self.conv(x)\n",
" y = self.bn(y)\n",
" y = self.relu(y)\n",
" return x + y\n",
"\n",
"model = Model()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The first step is to decide which part of the model should run with the quantized kernels. Since all the operations in this model support quantization, we may just quantize all inputs and dequantize all outputs. So we will get the modified model below."
]
},
{
"cell_type": "code",
"execution_count": 29,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
"\n",
"import torch.nn.quantized\n",
"\n",
"class Model(nn.Module):\n",
" def __init__(self):\n",
" super().__init__()\n",
" self.conv = nn.Conv2d(3, 3, 1)\n",
" self.bn = nn.BatchNorm2d(3)\n",
" self.relu = nn.ReLU()\n",
" self.fake_quant = torch.quantization.QuantStub() # Quantize\n",
" self.fake_dequant = torch.quantization.DeQuantStub() # Dequantize\n",
" \n",
" def forward(self, x):\n",
" x = self.fake_quant(x)\n",
" y = self.conv(x)\n",
" y = self.bn(y)\n",
" y = self.relu(y)\n",
"\n",
" z = x + y\n",
" z = self.fake_dequant(z)\n",
" return z"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The second step is to find out all the requantizable functions in the model. Wait, what does `requantizable` mean? It means the operations that may generate outputs with a different set of quantization parameters. Typically, the list include `add`, `mul`, `add_relu` and `cat`. We will need to replace them with the ones under `torch.nn.quantized.FloatFunctional`."
]
},
{
"cell_type": "code",
"execution_count": 30,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
"\n",
"import torch.nn.quantized\n",
"\n",
"class Model(nn.Module):\n",
" def __init__(self):\n",
" super().__init__()\n",
" self.conv = nn.Conv2d(3, 3, 1)\n",
" self.bn = nn.BatchNorm2d(3)\n",
" self.relu = nn.ReLU()\n",
" self.fake_quant = torch.quantization.QuantStub() # Quantize\n",
" self.fake_dequant = torch.quantization.DeQuantStub() # Dequantize\n",
" self.float_functional = torch.nn.quantized.FloatFunctional()\n",
" \n",
" def forward(self, x):\n",
" x = self.fake_quant(x)\n",
" y = self.conv(x)\n",
" y = self.bn(y)\n",
" y = self.relu(y)\n",
"\n",
" z = self.float_functional.add(x, y)\n",
" z = self.fake_dequant(z)\n",
" return z"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"With the model given above, you may use it in quantization. Next, we will need to figure out the fusable nodes. Some nodes can be viewed as one module during quantization, e.g. Conv2d-BatchNorm2d-ReLU."
]
},
{
"cell_type": "code",
"execution_count": 31,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Model(\n",
" (conv): ConvBnReLU2d(\n",
" (0): Conv2d(3, 3, kernel_size=(1, 1), stride=(1, 1))\n",
" (1): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (2): ReLU()\n",
" )\n",
" (bn): Identity()\n",
" (relu): Identity()\n",
" (fake_quant): QuantStub()\n",
" (fake_dequant): DeQuantStub()\n",
" (float_functional): FloatFunctional(\n",
" (activation_post_process): Identity()\n",
" )\n",
")\n"
]
}
],
"source": [
"from distutils.version import LooseVersion\n",
"\n",
"m = Model()\n",
"m.train()\n",
"\n",
"m.qconfig = torch.quantization.get_default_qat_qconfig('qnnpack')\n",
"\n",
"if LooseVersion(torch.__version__) >= LooseVersion('1.11.0'):\n",
" torch.ao.quantization.fuse_modules_qat(m, [['conv', 'bn', 'relu']], inplace=True)\n",
"else:\n",
" torch.quantization.fuse_modules(m, [['conv', 'bn', 'relu']], inplace=True)\n",
"\n",
"print(m)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The final step before model training or calibration is to perform quantization preparation. After this step, the `FakeQuantize` nodes will be added to all the requantizable nodes."
]
},
{
"cell_type": "code",
"execution_count": 32,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Model(\n",
" (conv): ConvBnReLU2d(\n",
" 3, 3, kernel_size=(1, 1), stride=(1, 1)\n",
" (bn): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (weight_fake_quant): FakeQuantize(\n",
" fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8), quant_min=-128, quant_max=127, dtype=torch.qint8, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])\n",
" (activation_post_process): MovingAverageMinMaxObserver(min_val=inf, max_val=-inf)\n",
" )\n",
" (activation_post_process): FakeQuantize(\n",
" fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8), quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])\n",
" (activation_post_process): MovingAverageMinMaxObserver(min_val=inf, max_val=-inf)\n",
" )\n",
" )\n",
" (bn): Identity()\n",
" (relu): Identity()\n",
" (fake_quant): QuantStub(\n",
" (activation_post_process): FakeQuantize(\n",
" fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8), quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])\n",
" (activation_post_process): MovingAverageMinMaxObserver(min_val=inf, max_val=-inf)\n",
" )\n",
" )\n",
" (fake_dequant): DeQuantStub()\n",
" (float_functional): FloatFunctional(\n",
" (activation_post_process): FakeQuantize(\n",
" fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8), quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1, scale=tensor([1.]), zero_point=tensor([0])\n",
" (activation_post_process): MovingAverageMinMaxObserver(min_val=inf, max_val=-inf)\n",
" )\n",
" )\n",
")\n"
]
}
],
"source": [
"torch.quantization.prepare_qat(m, inplace=True)\n",
"print(m)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, we may start our training process. To save time, we implemented the simple logic for feeding the model with some randomly-generated data."
]
},
{
"cell_type": "code",
"execution_count": 33,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Model(\n",
" (conv): ConvBnReLU2d(\n",
" 3, 3, kernel_size=(1, 1), stride=(1, 1)\n",
" (bn): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
" (weight_fake_quant): FakeQuantize(\n",
" fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8), quant_min=-128, quant_max=127, dtype=torch.qint8, qscheme=torch.per_tensor_symmetric, ch_axis=-1, scale=tensor([0.0045]), zero_point=tensor([0])\n",
" (activation_post_process): MovingAverageMinMaxObserver(min_val=-0.5689650177955627, max_val=0.4857633411884308)\n",
" )\n",
" (activation_post_process): FakeQuantize(\n",
" fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8), quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1, scale=tensor([0.0171]), zero_point=tensor([0])\n",
" (activation_post_process): MovingAverageMinMaxObserver(min_val=0.0, max_val=4.370081901550293)\n",
" )\n",
" )\n",
" (bn): Identity()\n",
" (relu): Identity()\n",
" (fake_quant): QuantStub(\n",
" (activation_post_process): FakeQuantize(\n",
" fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8), quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1, scale=tensor([0.0340]), zero_point=tensor([122])\n",
" (activation_post_process): MovingAverageMinMaxObserver(min_val=-4.156678199768066, max_val=4.5101141929626465)\n",
" )\n",
" )\n",
" (fake_dequant): DeQuantStub()\n",
" (float_functional): FloatFunctional(\n",
" (activation_post_process): FakeQuantize(\n",
" fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8), quant_min=0, quant_max=255, dtype=torch.quint8, qscheme=torch.per_tensor_affine, ch_axis=-1, scale=tensor([0.0460]), zero_point=tensor([87])\n",
" (activation_post_process): MovingAverageMinMaxObserver(min_val=-4.007132530212402, max_val=7.722269058227539)\n",
" )\n",
" )\n",
")\n"
]
}
],
"source": [
"for _ in range(10):\n",
" dummy_input = torch.randn(1, 3, 224, 224)\n",
" m(dummy_input)\n",
"\n",
"print(m)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As you see in the graph, the training process is carried out on the floating computation graph with the `FakeQuantize` nodes. So if you want a actual quantized model, we need to perform explicit conversion."
]
},
{
"cell_type": "code",
"execution_count": 34,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Model(\n",
" (conv): QuantizedConvReLU2d(3, 3, kernel_size=(1, 1), stride=(1, 1), scale=0.017137575894594193, zero_point=0)\n",
" (bn): Identity()\n",
" (relu): Identity()\n",
" (fake_quant): Quantize(scale=tensor([0.0340]), zero_point=tensor([122]), dtype=torch.quint8)\n",
" (fake_dequant): DeQuantize()\n",
" (float_functional): QFunctional(\n",
" scale=0.045997653156518936, zero_point=87\n",
" (activation_post_process): Identity()\n",
" )\n",
")\n"
]
}
],
"source": [
"quantized_m = torch.quantization.convert(m)\n",
"\n",
"print(quantized_m)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Our quantization tool\n",
"As you can see, a lot of things have to be done to apply quantization to your PyTorch model. That's why we develop the quantization tools in TinyNeuralNetwork, which eases the task by adding only several lines to your code."
]
},
{
"cell_type": "code",
"execution_count": 36,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Model_qat(\n",
" (fake_quant_0): Quantize(scale=tensor([0.0350]), zero_point=tensor([127]), dtype=torch.quint8)\n",
" (conv): QuantizedConvReLU2d(3, 3, kernel_size=(1, 1), stride=(1, 1), scale=0.017351238057017326, zero_point=0)\n",
" (bn): Identity()\n",
" (relu): Identity()\n",
" (fake_dequant_0): DeQuantize()\n",
" (float_functional_simple_0): QFunctional(\n",
" scale=0.03781222179532051, zero_point=94\n",
" (activation_post_process): Identity()\n",
" )\n",
")\n"
]
}
],
"source": [
"import sys\n",
"sys.path.append('../..')\n",
"\n",
"from tinynn.graph.quantization.quantizer import QATQuantizer\n",
"\n",
"quantizer = QATQuantizer(model, dummy_input, work_dir='out')\n",
"q_model = quantizer.quantize()\n",
"\n",
"for _ in range(10):\n",
" dummy_input = torch.randn(1, 3, 224, 224)\n",
" q_model(dummy_input)\n",
"\n",
"q_model = torch.quantization.convert(q_model)\n",
"\n",
"print(q_model)"
]
}
],
"metadata": {
"interpreter": {
"hash": "5a8cfc575211f63216cc03e2bf5e39a742bbf46e9fed10f94c831954dd3fbfef"
},
"kernelspec": {
"display_name": "Python 3.8.6 ('torch110': venv)",
"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.6"
},
"orig_nbformat": 4
},
"nbformat": 4,
"nbformat_minor": 2
}