release commit
[jacinto-ai/pytorch-jacinto-ai-devkit.git] / docs / Quantization.md
index 14c48810c2eb0a933fa63c176b8d0d0728497589..ff07b768fb6c4cef113c5b922e095efa8be5abd3 100644 (file)
@@ -1,43 +1,52 @@
 # Quantization
 
-As we know Quantization is the process of converting floating point data & operations to fixed point (integer). CNNs can be quantized to 8-bits integer data/operations without significant accuracy loss. This includes quantization of weights, feature maps and all operations (including convolution of weights). **We use power-of-2, symmetric quantization for both weights and activations**.
+Quantization is the process of converting floating point data & operations to fixed point (integer). CNNs can be quantized to 8-bits integer data/operations without significant accuracy loss. This includes quantization of weights, feature maps and all operations (including convolution of weights). The quantization style used in this code is **power-of-2, symmetric quantization for both weights and activations**.
 
-There are two primary methods of quantization - Post Training Quantization and Trained Quantization. 
+In order to make the activations quantization friendly, it is important to clip them during Calibration or Trained Quantization. PACT2 activation function has been developed to clip the activations to a power-of-two value. PACT2 can be used in the place of commonly used activation functions such as ReLU. 
+<p float="left"> <img src="quantization/pact2_activation.png" width="640" hspace="5"/> </p>
+This code also contains two wrapper modules called QuantCalibrateModule and QuantTrainModule that will handle such tasks as replacing ReLUs with PACT2, quantizing weights, quantizing activations etc - so the user does not need to do any of these things - the only thing that is required is to wrap the user's module in these wapper modules as explained below.
 
-## Post Training Calibration & Quantization
+There are two primary methods of quantization - Post Training Calibration For Quantization and Trained Quantization (a.k.a Quantization Aware Training). 
+
+## Post Training Calibration For Quantization
 
 Post Training Calibration & Quantization can take a model trained in floating point and with a few steps convert it to a model that is friendly for quantized inference. Compared to the alternative (Trained Quantization), the advantages of this method are:
 - Calibration is fast - a typical calibration finishes in a few minutes.  
 - Ground truth is not required - just input images are sufficient.
 - Loss function or backward (back-propagation) are not required. 
 
-Thus, this is the preferred method of quantization from an ease of use point of view.  As explained earlier, in this method, the training happens entirely in floating point. The inference (possibly in an embedded device) happens in fixed point. In between training and fixed point inference, the model goes through the step called Calibration with some sample images. The Calibration happens in PC and Quantized Inference happens in the embedded device. Calibration basically tries to make the quantized output similar to the floating point output - by choosing appropriate activation ranges, weights and biases. The step by step process is as follows:
+Thus, this is the preferred method of quantization from an ease of use point of view.  As explained earlier, in this method, the training happens entirely in floating point. The inference (possibly in an embedded device) happens in fixed point. In between training and fixed point inference, the model goes through the step called Calibration with some sample images. The Calibration happens in PC and Quantized Inference happens in the embedded device. Calibration basically tries to make the quantized output similar to the floating point output - by choosing appropriate activation ranges, weights and biases. 
 
-#### Model preparation:
-- Replace all the ReLU, ReLU6 layers in the model by PACT2. Insert PACT2 after Convolution+BatchNorm if a ReLU is missing after that.  Insert PACT2 anywhere else required - where activation range clipping and range collection is required. For example it can ne after the Fully Connected Layer. We use forward post hooks of PyTorch nn.Modules to call these extra activation functions. Thus we are able to add these extra activations without disturbing the loading of existing pre-trained weights.
-- Clip the weights to an appropriate range if the weight range is very high.
-
-#### Forward iterations:
-- For each iteration perform a forward in floating point using the original weights and biases. During this pass PACT2 layers will collect output ranges using histogram and running average.
-- In addition, perform Convolution+BatchNorm merging and quantization of the resulting weights. These quantized and de-quantized weights are used in a forward pass. Ranges collected by PACT2 is used for activation quantization (and de-quantization) to generate quantized output.
-- The floating point output and quantized output are compared using statistic measures. Using such statistic measures, we can adjust the weights and biases of Convolutions and Batch Normalization layers - so that the quantized output becomes closer to the floating point output.
-- Within a few iterations, we could get reasonable quantization accuracy for several models that we tried this method on.
+A block diagram of Post Training Calibration is shown below:
+<p float="left"> <img src="quantization/bias_calibration.png" width="640" hspace="5"/> </p>
 
 Depending on how the activation range is collected and Quantization is done, we have a few variants of this basic scheme.  
 - Simple Calib: Calibration includes PACT2 for activation clipping, running average and range collection. In this method we use min-max for activation range collection (no histogram).
 - **Advanced Calib**: Calibration includes PACT2 with histogram based ranges, Weight clipping, Bias correction. 
-- Advanced DW Calib: Calibration includes Per-Channel Quantization of Weights for Depthwise layers, PACT2 with histogram based ranges, Weight clipping, Bias correction. One of the earliest papers that clearly explained the benefits of Per-Channel Quantization for weights only (while the activations are quantized as Per-Tensor) is [6] 
+- Advanced DW Calib: Calibration includes Per-Channel Quantization of Weights for Depthwise layers, PACT2 with histogram based ranges, Weight clipping, Bias correction. One of the earliest papers that clearly explained the benefits of Per-Channel Quantization for weights only (while the activations are quantized as Per-Tensor) is [[6]]
 - Advanced Per-Chan Calib: Calibration includes Per-Channel Quantization for all layers, PACT2 with histogram based ranges, Weight clipping, Bias correction.
 
 Out of these methods, **Advanced Calib** is our recommended Calibration method as of now, as it has the best trade-off between the Accuracy and the features required during fixed point inference. All the Calibration scripts that we have in this page uses "Advanced Calib" by default. Other Calibration methods described here are for information only. 
 
-In order to do Calibration easily we have a developed a wrapper module called QuantCalibrateModule, which is located in pytorch_jacinto_ai.xnn.quantize.QuantCalibrateModule. We make use of a kind of Parametric Activation called **PACT2** in order to store the calibrated ranges of activations. PACT2 is a improved form of PACT [1]. **PACT2 uses power of 2 activation ranges** for activation clipping. PACT2 can learn ranges very quickly (using a statistic method) without back propagation - this feature makes it quite attractive for Calibration. Our wrapper module replaces all the ReLUs in the model with PACT2. It also inserts PACT2 in other places where activation ranges need to be collected.  Statistical range clipping in PACT2 improves the Quantized Accuracy over simple min-max range clipping. 
+#### How to use  QuantCalibrateModule
+- The section briefly explains how to make use of our helper/wrapper module to do the calibration of your model. For further details, please see pytorch_jacinto_ai.engine.train_classification.py and pytorch_jacinto_ai.engine.train_pixel2pixel.py. The step by step process is as follows:
+
+###### Model preparation:
+- Replace all the ReLU, ReLU6 layers in the model by PACT2. Insert PACT2 after Convolution+BatchNorm if a ReLU is missing after that.  Insert PACT2 anywhere else required - where activation range clipping and range collection is required. For example it can ne after the Fully Connected Layer. We use forward post hooks of PyTorch nn.Modules to call these extra activation functions. Thus we are able to add these extra activations without disturbing the loading of existing pre-trained weights.
+- Clip the weights to an appropriate range if the weight range is very high.
+- Note that this Model preparation is automatically done by QuantCalibrateModule
+
+###### Forward iterations:
+- For each iteration perform a forward in floating point using the original weights and biases. During this pass PACT2 layers will collect output ranges using histogram and running average.
+- In addition, perform Convolution+BatchNorm merging and quantization of the resulting weights. These quantized and de-quantized weights are used in a forward pass. Ranges collected by PACT2 is used for activation quantization (and de-quantization) to generate quantized output.
+- The floating point output and quantized output are compared using statistic measures. Using such statistic measures, we can adjust the weights and biases of Convolutions and Batch Normalization layers - so that the quantized output becomes closer to the floating point output.
+- Within a few iterations, we could get reasonable quantization accuracy for several models that we tried this method on.
 
-As explained, our method of **Calibration does not need ground truth, loss function or back propagation.** However in our script, we make use of ground truth to measure the loss/accuracy even in the Calibration stage - although that is not necessary
+In order to do Calibration easily we have a developed a wrapper module called QuantCalibrateModule, which is located in pytorch_jacinto_ai.xnn.quantize.QuantCalibrateModule. We make use of a kind of Parametric Activation called **PACT2** in order to store the calibrated ranges of activations. PACT2 is a improved form of PACT [[1]]. **PACT2 uses power of 2 activation ranges** for activation clipping. PACT2 can learn ranges very quickly (using a statistic method) without back propagation - this feature makes it quite attractive for Calibration. Our wrapper module replaces all the ReLUs in the model with PACT2. It also inserts PACT2 in other places where activation ranges need to be collected.  Statistical range clipping in PACT2 improves the Quantized Accuracy over simple min-max range clipping
 
-#### How to use  QuantCalibrateModule
-The section briefly explains how to make use of our helper/wrapper module to do the calibration of your model. For further details, please see pytorch_jacinto_ai.engine.train_classification.py and pytorch_jacinto_ai.engine.train_pixel2pixel.py.
+As explained, the method of **Calibration does not need ground truth, loss function or back propagation.** However in the calibration script, we make use of ground truth to measure the loss/accuracy even in the Calibration stage - although that is not necessary. 
 
+###### Sample Code Snippet:
 ```
 # create your model here:
 model = ...
@@ -67,36 +76,38 @@ torch.onnx.export(model.module, dummy_input, os.path.join(save_path,'model.onnx'
 
 Few examples of calibration are provided below. These commands are also listed in the file **run_quantization.sh** for convenience.<br>
 
-#### Calibration of ImageNet Classification MobileNetV2 model 
+- Calibration of ImageNet Classification MobileNetV2 model 
 ```
 python ./scripts/train_classification_main.py --phase calibration --dataset_name image_folder_classification --model_name mobilenetv2_tv_x1 --data_path ./data/datasets/image_folder_classification --pretrained https://download.pytorch.org/models/mobilenet_v2-b0353104.pth --batch_size 64 --quantize True --epochs 1 --epoch_size 100
 ```
 
-#### Calibration of ImageNet Classification ResNet50 model 
+- Calibration of ImageNet Classification ResNet50 model 
 ```
 python ./scripts/train_classification_main.py --phase calibration --dataset_name image_folder_classification --model_name mobilenetv2_tv_x1 --data_path ./data/datasets/image_folder_classification --pretrained https://download.pytorch.org/models/resnet50-19c8e357.pth --batch_size 64 --quantize True --epochs 1 --epoch_size 100
 ```
 
-#### Calibration of Cityscapes Semantic Segmentation model 
+- Calibration of Cityscapes Semantic Segmentation model 
 ```
 python ./scripts/train_segmentation_main.py --phase calibration --dataset_name cityscapes_segmentation --model_name deeplabv3lite_mobilenetv2_tv --data_path ./data/datasets/cityscapes/data --img_resize 384 768 --output_size 1024 2048 --gpus 0 1 
 --pretrained ./data/modelzoo/semantic_segmentation/cityscapes/deeplabv3lite-mobilenetv2/cityscapes_segmentation_deeplabv3lite-mobilenetv2_2019-06-26-08-59-32.pth 
 --batch_size 12 --quantize True --epochs 1 --epoch_size 100
 ```
 
-## Trained Quantization
-
+## Trained Quantization (a.k.a Quantization Aware Training)
 As explained in the previous section, Calibration is our preferred method of making a model quantization friendly. However, in exceptional cases, it is possible that the drop in accuracy during calibration is more than acceptable. In this case, Trained Quantization can be used. 
 
-Unlike Calibration, Trained Quantization involves ground truth, loss function and back propagation. The most popular method of trained quantization is [4]. It takes care of merging Convolution layers with the adjascent Batch Normalization layers (on-the-fly) during the quantized training (if this merging is not correctly done, quantized training may not improve the accuracy). In addition, we use Straight-Through Estimation (STE) [2,3] to improve the gradient flow in back-propagation. Also, the statistical range clipping in PACT2 improves the Quantized Accuracy over simple min-max range clipping. 
+Unlike Calibration, Trained Quantization involves ground truth, loss function and back propagation. The most popular method of trained quantization is [[4]]. It takes care of merging Convolution layers with the adjascent Batch Normalization layers (on-the-fly) during the quantized training (if this merging is not correctly done, quantized training may not improve the accuracy). In addition, we use Straight-Through Estimation (STE) [[2,3]] to improve the gradient flow in back-propagation. Also, the statistical range clipping in PACT2 improves the Quantized Accuracy over simple min-max range clipping. 
+
+Note: Instead of STE and statistical ranges for PACT2, we also tried out approximate gradients for scale and trained quantization thresholds proposed in [[5]] (We did not use the gradient nomralization and log-domain training mentioned in the paper). We found that method to be able to learn the clipping thresholds for initial few epochs, but became unstable after a few epochs and loss became high. Compared to that learned thresholds method, our statistical PACT2 ranges/thresholds combined with STE is simple and stable. 
 
-Note: Instead of STE and statistical ranges for PACT2, we also tried out approximate gradients for scale and trained quantization thresholds proposed in [5] (We did not use the gradient nomralization and log-domain training mentioned in the paper). We found that method to be able to learn the clipping thresholds for initial few epochs, but became unstable after a few epochs and loss became high. Compared to that learned thresholds method, our statistical PACT2 ranges/thresholds combined with STE is simple and stable. 
+A block diagram of Trained Quantization is shown below:
+<p float="left"> <img src="quantization/trained_quant_ste.png" width="640" hspace="5"/> </p>
 
 In order to enable quantized training, we have developed the wrapper class pytorch_jacinto_ai.xnn.quantize.QuantTrainModule. The usage of this module can be seen in pytorch_jacinto_ai.engine.train_classification.py and pytorch_jacinto_ai.engine.train_pixel2pixel.py. 
 ```
 model = pytorch_jacinto_ai.xnn.quantize.QuantTrainModule(model, dummy_input=dummy_input)
 ```
-The resultant model can then be used for training as usual and it will take care of quantization constraints during the training forward and backward passes.
+The Model Preparation steps used for Calibration applies for Trained Quantization as well. - Note that this Model preparation is automatically done by QuantTrainModule. The resultant model can then be used for training as usual and it will take care of quantization constraints during the training forward and backward passes.
 
 One word of caution is that our current implementation of Trained Quantization is a bit slow. The reason for this slowdown is that our implementation is using the top-level python layer of PyTorch and not the underlying C++ layer. But with PyTorch natively supporting the functionality required for quantization under the hood - we hope that this speed issue can be resolved in a future update. 
 
@@ -115,7 +126,7 @@ python ./scripts/train_segmentation_main.py --dataset_name cityscapes_segmentati
 - This may not be such a problem as calibration and quantization may not take as much time as the original training. 
 - If your calibration/training crashes with insufficient GPU memory, reduce the batch size and try again.
 - The original training (without quantization) can use Multi-GPU as usual and we do not have any restrictions on that.
-- Tools for Calibration and Trained Quantization have started appearing in mainstream Deep Learning training frameworks [7,8]. Using the tools natively provided by these frameworks may be faster compared to an implementation in the Python layer of these frameworks (like we have done) - but they may not be mature currently. 
+- Tools for Calibration and Trained Quantization have started appearing in mainstream Deep Learning training frameworks [[7,8]]. Using the tools natively provided by these frameworks may be faster compared to an implementation in the Python layer of these frameworks (like we have done) - but they may not be mature currently. 
 - In order that the activation range estimation is correct, **the same module should not be re-used multiple times within the module**. Unfortunately, in the torchvision ResNet models, the ReLU module in the BasicBlock and BottleneckBlock are re-used multiple times. We have corrected this by defining separate ReLU modules. This change is minor and **does not** affect the loading of existing pretrained weights. See the [our modified ResNet model definition here](./modules/pytorch_jacinto_ai/vision/models/resnet.py).
 - Use Modules instead of functions as much as possible (we make use of modules to decide whether to do activation range clipping or not). For example use torch.nn.AdaptiveAvgPool2d() instead of torch.nn.functional.adaptive_avg_pool2d(), torch.nn.Flatten() instead of torch.nn.functional.flatten() etc. If you are using functions in your model and is giving poor quantized accuracy, then consider replacing those functions by the corresponding modules.