quantization doc update
authorManu Mathew <a0393608@ti.com>
Fri, 9 Oct 2020 08:25:29 +0000 (13:55 +0530)
committerManu Mathew <a0393608@ti.com>
Fri, 9 Oct 2020 08:48:39 +0000 (14:18 +0530)
quantization doc update

doc update

README.md
docs/Quantization.md

index 97e0c1faafce20f0e5a87105f318cdc39d5ff402..5a000b3c7f417d01886a376d35889240c163f8b1 100644 (file)
--- a/README.md
+++ b/README.md
@@ -1,49 +1,89 @@
 # Jacinto-AI-DevKit (PyTorch)
+Training & Quantization Tools For Embedded AI Development - in PyTorch.
 
-###### Notice: 
+
+## Notice
 - If you have not visited our landing page in github, please do so: [https://github.com/TexasInstruments/jacinto-ai-devkit](https://github.com/TexasInstruments/jacinto-ai-devkit)
 - **Issue Tracker for jacinto-ai-devkit:** You can file issues or ask questions at **e2e**: [https://e2e.ti.com/support/processors/f/791/tags/jacinto_2D00_ai_2D00_devkit](https://e2e.ti.com/support/processors/f/791/tags/jacinto_2D00_ai_2D00_devkit). While creating a new issue kindly include **jacinto-ai-devkit** in the tags (as you create a new issue, there is a space to enter tags, at the bottom of the page). 
 - **Issue Tracker for TIDL:** [https://e2e.ti.com/support/processors/f/791/tags/TIDL](https://e2e.ti.com/support/processors/f/791/tags/TIDL). Please include the tag **TIDL** (as you create a new issue, there is a space to enter tags, at the bottom of the page). 
 - If you do not get a reply within two days, please contact us at: jacinto-ai-devkit@list.ti.com
 
-### Deep Learning Models / Training / Calibration & Quantization - Using PyTorch<br>
+<hr>
+
+
+## Introduction
 This code provides a set of low complexity deep learning examples and models for low power embedded systems. Low power embedded systems often requires balancing of complexity and accuracy. This is a tough task and requires significant amount of expertise and experimentation. We call this process **complexity optimization**. In addition we would like to bridge the gap between Deep Learning training frameworks and real-time embedded inference by providing ready to use examples and enable **ease of use**. Scripts for training, validation, complexity analysis are also provided. 
 
 This code also includes tools for **Quantization Aware Training** that can output an 8-bit Quantization friendly model - these tools can be used to improve the quantized accuracy and bring it near floating point accuracy. For more details, please refer to the section on [Quantization](docs/Quantization.md).
 
 **Several of these models have been verified to work on [TI's Jacinto7 Automotive Processors](http://www.ti.com/processors/automotive-processors/tdax-adas-socs/overview.html).** These tools and software are primarily intended as examples for learning and research.
 
+<hr>
+
+
 ## Installation Instructions
-- These instructions are for installation on **Ubuntu 18.04**. 
-- Install Anaconda with Python 3.7 or higher from https://www.anaconda.com/distribution/ <br>
-- After installation, make sure that your python is indeed Anaconda Python 3.7 or higher by typing:<br>
-    ```
-    python --version
-    ```
-- Clone this repository into your local folder
-- Execute the following shell script to install the dependencies:<br>
-    ```
-    ./setup.sh
-    ```
+These instructions are for installation on **Ubuntu 18.04**. 
+
+Install Anaconda with Python 3.7 or higher from https://www.anaconda.com/distribution/ <br>
+
+After installation, make sure that your python is indeed Anaconda Python 3.7 or higher by typing:<br>
+```
+python --version
+```
+
+Clone this repository into your local folder
+
+Execute the following shell script to install the dependencies:<br>
+```
+./setup.sh
+```
+
+<hr>
+
 
 ## Examples
-- [**Image Classification**](docs/Image_Classification.md)<br>
-- [**Semantic Segmentation**](docs/Semantic_Segmentation.md)<br>
-- [Depth Estimation](docs/Depth_Estimation.md)<br>
-- [Motion Segmentation](docs/Motion_Segmentation.md)<br>
-- [**Multi Task Estimation**](docs/Multi_Task_Learning.md)<br>
-- Object Detection - coming soon..<br>
-- Object Keypoint Estimation - coming soon..<br>
-- [**Quantization**](docs/Quantization.md)<br>
+Below are some of the examples are currently available. Click on each of the links above to go into the full description of the example.
+
+[**Image Classification**](docs/Image_Classification.md)<br>
+
+[**Semantic Segmentation**](docs/Semantic_Segmentation.md)<br>
+
+[**Object Detection**](https://git.ti.com/cgit/jacinto-ai/pytorch-mmdetection/about/) - this link will take you to another repository, where we have our object detection training scripts.
+
+[Depth Estimation](docs/Depth_Estimation.md)<br>
+
+[Motion Segmentation](docs/Motion_Segmentation.md)<br>
 
-Above are some of the examples are currently available. Click on each of the links above to go into the full description of the example.
+[**Multi Task Estimation**](docs/Multi_Task_Learning.md)<br>
+
+Object Keypoint Estimation - coming soon..<br>
+
+[**Quantization**](docs/Quantization.md)<br>
+
+<hr>
+
+
+## Model Quantization
+Quantization (especially 8-bit Quantization) is important to get best throughput for inference. Quantization can be done using either **Post Training Quantization (PTQ)** or **Quantization Aware Training (QAT)**.
+
+[TI Deep Learning Library (TIDL)](https://software-dl.ti.com/jacinto7/esd/processor-sdk-rtos-jacinto7/latest/exports/docs/psdk_rtos_auto/docs/user_guide/sdk_components.html#ti-deep-learning-library-tidl) that is part of the [Processor SDK RTOS for Jacinto7](https://software-dl.ti.com/jacinto7/esd/processor-sdk-rtos-jacinto7/latest/exports/docs/psdk_rtos_auto/docs/user_guide/index.html) natively supports **PTQ** - TIDL can take floating point models and can quantize them using advanced calibration methods. 
+
+We have  guidelines on how to choose models and how train them to get best accuracy with Quantization. It is unlikely that there will be significant accuracy drop with **PTQ** if these guidelines are followed. In spite of this, if there are models that have significant accuracy drop with quantization, it is possible to improve the accuracy using **QAT**. Please read more details in the documentation on **[Quantization](docs/Quantization.md)**.
+
+<hr>
 
 
 ## Additional Information
-- Some of the common training and validation commands are provided in shell scripts (.sh files) in the root folder. <br>
-- Landing Page: [https://github.com/TexasInstruments/jacinto-ai-devkit](https://github.com/TexasInstruments/jacinto-ai-devkit) <br>
-- Actual Git Repositories: [https://git.ti.com/jacinto-ai](https://git.ti.com/jacinto-ai) <br>
-- Each of the repositories listed in the above link have an "about" tab with documentation and a "summary" tab with git clone/pull URLs. 
+Some of the common training and validation commands are provided in shell scripts (.sh files) in the root folder. <br>
+
+Landing Page: [https://github.com/TexasInstruments/jacinto-ai-devkit](https://github.com/TexasInstruments/jacinto-ai-devkit) <br>
+
+Actual Git Repositories: [https://git.ti.com/jacinto-ai](https://git.ti.com/jacinto-ai) <br>
+
+Each of the repositories listed in the above link have an "about" tab with documentation and a "summary" tab with git clone/pull URLs. 
+
+<hr>
+
 
 ## Acknowledgements
 Our source code uses parts of the following open source projects. We would like to sincerely thank their authors for making their code bases publicly available.
@@ -54,6 +94,9 @@ Our source code uses parts of the following open source projects. We would like
 |Training, Validation Engine/Loops |https://github.com/pytorch/examples, https://github.com/ClementPinard/FlowNetPytorch |
 |Object Detection                  |https://github.com/open-mmlab/mmdetection                                            |
 
+<hr>
+
+
 ## License
 
 Please see the [LICENSE](./LICENSE) file for more information about the license under which this code is made available.
\ No newline at end of file
index 7e03c52601d900a8a9b9a9af975ebf9ea7ab3c1a..85aa23c36f3e311b03c9a96658f1255a66b03ea3 100644 (file)
@@ -4,74 +4,107 @@ Quantization of a CNN model is the process of converting floating point data & o
 
 Accuracy of inference can degrade if the CNN model is quantized to 8bits using simple methods and steps have to be taken to minimize this accuracy loss. The parameters of the model need to be adjusted to suit quantization. This includes adjusting of weights, biases and activation ranges. This adjustment can be done as part of the Calibration or as part of Quantization Aware Training.
 
+
 ## Overview
-- Inference engines use fixed point arithmetic to implement neural networks. For example TI Deep Learning Library (TIDL) for TI’s Jacinto7 TDA4x Devices (eg. TDA4VM) supports 16-bit and 8-bit fixed point inference modes.
-- Fixed point mode, especially the 8-bit mode can have accuracy degradation. The tools and guidelines provided here help to avoid accuracy degradation with quantization.
-- If you are getting accuracy degradation with 8-bit inference, the first thing to check is 16-bit inference. If 16-bit inference provides accuracy close to floating point and 8-bit has an accuracy degradation, there it is likely that the degradation si due to quantization. However, if there is substantial accuracy degradation with 16-bit inference itself, then it is likely that there there is some issue other than quantization.  
+Inference engines use fixed point arithmetic to implement neural networks. For example TI Deep Learning Library (TIDL) for TI’s Jacinto7 TDA4x Devices (eg. TDA4VM) supports 16-bit and 8-bit fixed point inference modes.
+
+Fixed point mode, especially the 8-bit mode can have accuracy degradation. The tools and guidelines provided here help to avoid accuracy degradation with quantization.
+
+If you are getting accuracy degradation with 8-bit inference, the first thing to check is 16-bit inference. If 16-bit inference provides accuracy close to floating point and 8-bit has an accuracy degradation, there it is likely that the degradation si due to quantization. However, if there is substantial accuracy degradation with 16-bit inference itself, then it is likely that there there is some issue other than quantization.  
+
 
 #### Quantization Schemes
-- Post Training Calibration & Quantization (Calibration): Calibration often involves range estimation for weights and activations and also minor tweaks to the model (such as bias adjustments). Fixed point inference engines such as TIDL can accept a floating point model and Calibrate it using a few sample images. The Calibration is done during the import of the model, in the case of TIDL.<br>
-- Quantization Aware Training (QAT): This is needed if accuracy obtained with Calibration is not satisfactory (eg. Quantization Accuracy Drop >2%). QAT operates as a second phase after the initial training in floating point is done. 
+Post Training Calibration & Quantization (Calibration): Calibration often involves range estimation for weights and activations and also minor tweaks to the model (such as bias adjustments). Fixed point inference engines such as TIDL can accept a floating point model and Calibrate it using a few sample images. The Calibration is done during the import of the model, in the case of TIDL.<br>
+
+Quantization Aware Training (QAT): This is needed if accuracy obtained with Calibration is not satisfactory (eg. Quantization Accuracy Drop >2%). QAT operates as a second phase after the initial training in floating point is done. 
+
 
 ## Guidelines For Training To Get Best Accuracy With Quantization
-- **These are important** - we are listing these guidelines upfront because it is important to follow these.
-- We recommend that the training uses sufficient amount of regularization / weight decay. Regularization / weight decay ensures that the weights, biases and other parameters (if any) are small and compact - this is good for quantization. These features are supported in most of the popular training framework.<br>
-- We have noticed that some training code bases do not use weight decay for biases. Some other code bases do not use weight decay for the parameters in Depthwise convolution layers. All these are bad strategies for quantization. These poor choices done (probably to get a 0.1% accuracy lift with floating point) will result in a huge degradation in fixed point - sometimes several percentage points. The weight decay factor should not be too small. We have used a weight decay factor of 1e-4 for training several networks and we highly recommend a similar value. Please do no use small values such as 1e-5.<br>
-- We also highly recommend to use Batch Normalization immediately after every Convolution layer. This helps the feature map to be properly regularized/normalized. If this is not done, there can be accuracy degradation with quantization. This especially true for Depthwise Convolution layers. However applying Batch Normalization to the very last Convolution layer (for example, the prediction layer in segmentation/object detection network) may hurt accuracy and can be avoided.<br>
+**These are important** - we are listing these guidelines upfront because it is important to follow these.
+
+We recommend that the training uses **sufficient amount of regularization (weight decay) for all parameters**. Regularization / weight decay ensures that the weights, biases and other parameters (if any) are small and compact - this is good for quantization. These features are supported in most of the popular training framework.<br>
+
+We have noticed that some training code bases do not use weight decay for biases. Some other code bases do not use weight decay for the parameters in Depthwise convolution layers. All these are bad strategies for quantization. These poor choices done (probably to get a 0.1% accuracy lift with floating point) will result in a huge degradation in fixed point - sometimes several percentage points. The weight decay factor should not be too small. We have used a weight decay factor of 1e-4 for training several networks and we highly recommend a similar value. Please do no use small values such as 1e-5.<br>
+
+We also highly recommend to use **Batch Normalization immediately after every Convolution layer**. This helps the feature map to be properly regularized/normalized. If this is not done, there can be accuracy degradation with quantization. This especially true for Depthwise Convolution layers. However applying Batch Normalization to the Convolution layer that does the prediction (for example, the very last layer in segmentation/object detection network) may hurt accuracy and can be avoided.<br>
 
-To get best accuracy at the quantization stage, it is important that the model is trained carefully, following the guidelines (even during the floating point training). Having spent a lot of time solving quantization issues, we would like to highlight that following these guidelines are of at most importance. Otherwise, there is a high risk that the tools and techniques described here may not be completely effective in solving the accuracy drop due to quantization. To summarize, if you are getting poor accuracy with quntization, please check the following:<br>
+There are several model types out there including MobileNets, ResNets, DenseNets, EfficientNets(Lite), RegNetX [9] etc. that are popular in the embedded community. The models using Depthwise convolutions (such as MobileNets) are more difficult to quantize - expect higher accuracy drop with quantization when using such models. We would like to **recommend RegNetX [9] models** as the most embedded friendly as they balance accuracy for a given complexity and the ease of quantization due to their use of grouped convolutions with carefully selected group sizes.<br>
+
+To get best accuracy at the quantization stage, it is important that the model is trained carefully, following the guidelines (even during the floating point training). Having spent a lot of time solving quantization issues, we would like to highlight that following these guidelines are of at most importance. Otherwise, there is a high risk that the tools and techniques described here may not be completely effective in solving the accuracy drop due to quantization. To summarize, if you are getting poor accuracy with quantization, please check the following:<br>
 - Weight decay is applied to all layers / parameters and that weight decay factor is good.<br>
 - Ensure that the Convolution layers in the network have Batch Normalization layers immediately after that. The only exception allowed to this rule is for the very last Convolution layer in the network (for example the prediction layer in a segmentation network or detection network, where adding Batch normalization might hurt the floating point accuracy).<br>
 
+
 ## Implementation Notes, Limitations & Recommendations
-- **Please read carefully** - closely following these recommendations can save hours or days of debug related to quantization accuracy issues.
-- **The same module should not be re-used multiple times within the module** in order that the feature map range estimation is correct. 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/xvision/models/resnet.py).<br>
-- **Use Modules instead of functionals or tensor operations** (by Module we mean classes derived from torch.nn.Module). We make use of Modules heavily in our quantization tools - in order to do range collection, in order to merge Convolution/BatchNorm/ReLU in order to decide whether to quantize a certain tensor and so on. For example use torch.nn.ReLU instead of torch.nn.functional.relu(), torch.nn.AdaptiveAvgPool2d() instead of torch.nn.functional.adaptive_avg_pool2d(), torch.nn.Flatten() instead of torch.nn.functional.flatten() etc.<br>
-- Other notable modules provided are: [xnn.layers.AddBlock](../modules/pytorch_jacinto_ai/xnn/layers/common_blocks.py) to do elementwise addition and [xnn.layers.CatBlock](../modules/pytorch_jacinto_ai/xnn/layers/common_blocks.py) to do concatenation of tensors. Use these in the models instead of tensor operations. Note that if there are multiple element wise additions in a model, each of them should use a different instance of xnn.layers.AddBlock (since the same module should not be re-used multiple times - see above). The same restriction applies for xnn.layers.CatBlock or any other module as well.
-- **Interpolation/Upsample/Resize** has been tricky in PyTorch in the sense that the ONNX graph generated used to be unnecessarily complicated. Recent versions of PyTorch has fixed it - but the right options must be used to get the clean graph. We have provided a functional form as well as a module form of this operator with the capability to export a clean ONNX graph [xnn.layers.resize_with, xnn.layers.ResizeWith](../modules/pytorch_jacinto_ai/xnn/layers/resize_blocks.py)
-- If you have done QAT and is getting poor accuracy either in the Python code or during inference in the platform, please inspect your model carefully to see if the above recommendations have been followed - some of these can be easily missed by oversight - and can result in painful debugging that could have been avoided.<br>
-- However, if a function does not change the range of feature map, it is not critical to use it in Module form. An example of this is torch.nn.functional.interpolate<br>
-- **Multi-GPU training/validation with DataParallel** is supported with our QAT module QuantTrainModule and Test module QuantTestModule. This takes care of a major concern that was earlier there in doing QAT with QuantTrainModule. (However it is not supported for QuantCalibrateModule - calibration take much less time - so hopefully this is not a big issue. In our example training scripts train_classification.py and train_pixel2pixel.py in pytorch_jacinto_ai/engine, we do not wrap the model in DataParallel if the model is QuantCalibrateModule, but we do that for QuantTrainModule and QuantTestModule).<br>
-- If your training/calibration crashes because of insufficient GPU memory, reduce the batch size and try again.
-- If you are using TIDL to infer a model trained using QAT (or Calibrated model using the PTQ Calibration that is simulated here) tools provided in this repository, please set the following in the import config file of TIDL for best accuracy: <br>
+**Please read carefully** - closely following these recommendations can save hours or days of debug related to quantization accuracy issues.
+
+**The same module should not be re-used multiple times within the module** in order that the feature map range estimation is correct. 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/xvision/models/resnet.py).<br>
+
+**Use Modules instead of functionals or tensor operations** (by Module we mean classes derived from torch.nn.Module). We make use of Modules heavily in our quantization tools - in order to do range collection, in order to merge Convolution/BatchNorm/ReLU in order to decide whether to quantize a certain tensor and so on. For example use torch.nn.ReLU instead of torch.nn.functional.relu(), torch.nn.AdaptiveAvgPool2d() instead of torch.nn.functional.adaptive_avg_pool2d(), torch.nn.Flatten() instead of torch.nn.functional.flatten() etc.<br>
+
+Other notable modules provided are: [xnn.layers.AddBlock](../modules/pytorch_jacinto_ai/xnn/layers/common_blocks.py) to do elementwise addition and [xnn.layers.CatBlock](../modules/pytorch_jacinto_ai/xnn/layers/common_blocks.py) to do concatenation of tensors. Use these in the models instead of tensor operations. Note that if there are multiple element wise additions in a model, each of them should use a different instance of xnn.layers.AddBlock (since the same module should not be re-used multiple times - see above). The same restriction applies for xnn.layers.CatBlock or any other module as well.
+
+**Interpolation/Upsample/Resize** has been tricky in PyTorch in the sense that the ONNX graph generated used to be unnecessarily complicated. Recent versions of PyTorch has fixed it - but the right options must be used to get the clean graph. We have provided a functional form as well as a module form of this operator with the capability to export a clean ONNX graph [xnn.layers.resize_with, xnn.layers.ResizeWith](../modules/pytorch_jacinto_ai/xnn/layers/resize_blocks.py)
+
+If you have done QAT and is getting poor accuracy either in the Python code or during inference in the platform, please inspect your model carefully to see if the above recommendations have been followed - some of these can be easily missed by oversight - and can result in painful debugging that could have been avoided.<br>
+
+However, if a function does not change the range of feature map, it is not critical to use it in Module form. An example of this is torch.nn.functional.interpolate<br>
+
+**Multi-GPU training/validation with DataParallel** is supported with our QAT module QuantTrainModule and Test module QuantTestModule. This takes care of a major concern that was earlier there in doing QAT with QuantTrainModule. (However it is not supported for QuantCalibrateModule - calibration take much less time - so hopefully this is not a big issue. In our example training scripts train_classification.py and train_pixel2pixel.py in pytorch_jacinto_ai/engine, we do not wrap the model in DataParallel if the model is QuantCalibrateModule, but we do that for QuantTrainModule and QuantTestModule).<br>
+
+If your training/calibration crashes because of insufficient GPU memory, reduce the batch size and try again.
+
+If you are using TIDL to infer a model trained using QAT (or Calibrated model using the PTQ Calibration that is simulated here) tools provided in this repository, please set the following in the import config file of TIDL for best accuracy: <br>
   **quantizationStyle = 3** to use power of 2 quantization. <br> 
   **foldPreBnConv2D = 0** to avoid a slight accuracy degradation due to incorrect folding of BatchNormalization that comes before Convolution (input mean/scale is implemented in TIDL as a PreBN - so this affects most networks). <br> 
   **calibrationOption = 0** to avoid further Calibration in TIDL. <br>
 
+
 ## Post Training Calibration For Quantization (PTQ a.k.a. Calibration)
-- **Note: this is not our recommended method in PyTorch.**<br>
-- Post Training Calibration or simply Calibration is a method to reduce the accuracy loss with quantization. This is an approximate method and does not require ground truth or back-propagation - hence it is suitable for implementation in an Import/Calibration tool. 
-- For example, PTQ with Advanced Calibration can be enabled in TIDL by setting **calibrationOption = 7**. Please consult the TIDL documentation for further explanation fo this option.
-- We have simulated PTQ with Advanced Calibration in PyTorch. If you are interested, you can take a look at the [documentation of Calibration here](Calibration.md).<br>
-- However, in a training frame work such as PyTorch, it is possible to get better accuracy with Quantization Aware Training (QAT) and we recommend to use that (next section).
+**Note: this is not our recommended method in PyTorch.**<br>
+
+Post Training Calibration or simply Calibration is a method to reduce the accuracy loss with quantization. This is an approximate method and does not require ground truth or back-propagation - hence it is suitable for implementation in an Import/Calibration tool. 
+
+For example, PTQ with Advanced Calibration can be enabled in TIDL by setting **calibrationOption = 7**. Please consult the TIDL documentation for further explanation fo this option.
+
+We have simulated PTQ with Advanced Calibration in PyTorch. If you are interested, you can take a look at the [documentation of Calibration here](Calibration.md).<br>
+
+However, in a training frame work such as PyTorch, it is possible to get better accuracy with Quantization Aware Training (QAT) and we recommend to use that (next section).
+
 
 ## Quantization Aware Training (QAT)
 Quantization Aware Training (QAT) is easy to incorporate into an existing PyTorch training code. We provide a wrapper module called QuantTrainModule to automate all the tasks required for QAT. The user simply needs to wrap his model in QuantTrainModule and do the training.
 
-The overall flow of training is as follows:
-- Step 1:Train your model in floating point as usual.
-- Step 2: Starting from the floating point model as pretrained weights, do Quantization Aware Training. In order to do this wrap your model in the wrapper module called  pytorch_jacinto_ai.xnn.quantize.QuantTrainModule and perform training with a small learning rate. About 25 to 50 epochs of training may be required to get the best accuracy.
+The overall flow of training is as follows:<br>
+- Step 1:Train your model in floating point as usual.<br>
+- Step 2: Starting from the floating point model as pretrained weights, do Quantization Aware Training. In order to do this wrap your model in the wrapper module called  pytorch_jacinto_ai.xnn.quantize.QuantTrainModule and perform training with a small learning rate. About 25 to 50 epochs of training may be required to get the best accuracy.<br>
 
-QuantTrainModule does the following operations to the model. Note that QuantTrainModule that will handle these tasks - the only thing that is required is to wrap the user's module in QuantTrainModule as explained in the section "How to use  QuantTrainModule".
-- Replace all the ReLU, ReLU6 layers in the model by PACT2. Insert PACT2 after Convolution+BatchNorm if a ReLU/ReLU6 is missing after that.  Insert PACT2 anywhere else required - where activation range clipping and range collection is required. For example it can be 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.
-- Quantize the weights during the forward pass. Merging Convolution layers with the adjacent Batch Normalization layers (on-the-fly) during the weight quantization is required - if this merging is not correctly done, Quantization Aware Training may not improve accuracy.
-- Quantize activations during the forward pass.
-- Other modifications to help the learning process. For example, we use Straight-Through Estimation (STE) [[2,3]] to improve the gradient flow in back-propagation.
+QuantTrainModule does the following operations to the model. Note that QuantTrainModule that will handle these tasks - the only thing that is required is to wrap the user's module in QuantTrainModule as explained in the section "How to use  QuantTrainModule".<br>
+- Replace all the ReLU, ReLU6 layers in the model by PACT2. Insert PACT2 after Convolution+BatchNorm if a ReLU/ReLU6 is missing after that.  Insert PACT2 anywhere else required - where activation range clipping and range collection is required. For example it can be 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.<br>
+- Clip the weights to an appropriate range if the weight range is very high.<br>
+- Quantize the weights during the forward pass. Merging Convolution layers with the adjacent Batch Normalization layers (on-the-fly) during the weight quantization is required - if this merging is not correctly done, Quantization Aware Training may not improve accuracy.<br>
+- Quantize activations during the forward pass.<br>
+- Other modifications to help the learning process. For example, we use Straight-Through Estimation (STE) [[2,3]] to improve the gradient flow in back-propagation.<br>
 
 A block diagram of Quantization Aware Training with QuantTrainModule is shown below:
 <p float="left"> <img src="quantization/trained_quant_ste.png" width="640" hspace="5"/> </p>
 
+
 #### PACT2 activation
 In order to make the activations quantization friendly, it is important to clip them during Quantization Aware Training. PACT2 activation module has been developed to clip the activations to a power-of-two value. PACT2 is used in the place of commonly used activation functions such as ReLU or ReLU6. Our Quantization Aware Training modules/scripts will automatically insert PACT2 activation functions wherever necessary to constraint the ranges of activations. The following is a block diagram of the PACT2:
 <p float="left"> <img src="quantization/pact2_activation.png" width="640" hspace="5"/> </p>
 We use statistical range clipping in PACT2 to improve the Quantized Accuracy (compared to simple min-max range clipping).
 
+
 #### What happens during Quantization Aware Training?
-- 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.
-- Back-propagation with STE will update the parameters of the model to reduce the loss with quantization.
-- Within a few epochs, we should get reasonable quantization accuracy.
+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.
+
+Back-propagation with STE will update the parameters of the model to reduce the loss with quantization.
+
+Within a few epochs, we should get reasonable quantization accuracy.
+
 
 #### How to use  QuantTrainModule
 In order to enable quantized training, we have developed the wrapper class pytorch_jacinto_ai.xnn.quantize.QuantTrainModule. A simple example for using this module is given in the script [examples/quantization_example.py](../examples/quantization_example.py) and calling this is demonstrated in [run_quantization_example.sh](../run_quantization_example.sh). The usage of this module can also be seen in pytorch_jacinto_ai.engine.train_classification.py and pytorch_jacinto_ai.engine.train_pixel2pixel.py. The following is a brief description of how to use this wrapper module:
@@ -110,6 +143,7 @@ As can be seen, it is easy to incorporate QuantTrainModule in your existing trai
 
 Optional: We have provided a utility function called pytorch_jacinto_ai.xnn.utils.load_weights() that prints which parameters are loaded correctly and which are not - you can use this load function if needed to ensure that your parameters are loaded correctly.
 
+
 ####  Example commands for QAT
 ImageNet Classification: *In this example, only a fraction of the training samples are used in each training epoch to speedup training. Remove the argument --epoch_size to use all the training samples.*
 ```
@@ -123,6 +157,7 @@ python ./scripts/train_segmentation_main.py --dataset_name cityscapes_segmentati
 
 For more examples, please see the files run_qunatization_example.sh and examples/quantization_example.py
 
+
 ## Results
 
 The table below shows the Quantized Accuracy with various Calibration and methods and also QAT. Some of the commands used to generate these results are summarized in the file **run_quantization.sh** for convenience.
@@ -135,10 +170,11 @@ The table below shows the Quantized Accuracy with various Calibration and method
 |MobileNetV2(TV)    |MobileNetV2|32    |224x224   |**71.89** |67.77            |**68.39**     |69.34            |**70.74**|-3.50             |-1.34        |
 |MobileNetV2(Shicai)|MobileNetV2|32    |224x224   |**71.44** |0.0              |**68.81**     |70.65            |**70.54**|-2.63             |-0.9         |
 
-Notes:
+Notes:<br>
 - For Image Classification, the accuracy measure used is % Top-1 Classification Accuracy. 'Top-1 Classification Accuracy' is abbreviated by Acc in the above table.<br>
 - (TV) Stands for TochVision: https://github.com/pytorch/vision
-- MobileNetV2(Shicai) model is from https://github.com/shicai/MobileNet-Caffe (converted from caffe to PyTorch) - this model was selected as this is a tough case for quantization.<br><br>
+- MobileNetV2(Shicai) model is from https://github.com/shicai/MobileNet-Caffe (converted from caffe to PyTorch) - this model was selected as this is a tough case for quantization.<br>
+
 
 ###### Dataset: Cityscapes Segmentation (Semantic Segmentation)
 
@@ -146,18 +182,19 @@ Notes:
 |----------   |-----------|------|----------|----------|---              |---           |---              |---      |---                    |---           |
 |DeepLabV3Lite|MobileNetV2|16    |768x384   |**69.13** |61.71            |**67.95**     |68.47            |**68.44**|-1.18                  |-0.69         |
 
-Note: For Semantic Segmentation, the accuracy measure used in MeanIoU Accuracy. 'MeanIoU Accuracy' is abbreviated by Acc in the above table.
+Note:<br>
+- For Semantic Segmentation, the accuracy measure used in MeanIoU Accuracy. 'MeanIoU Accuracy' is abbreviated by Acc in the above table.
 
 **Terminology:**<br>
 All of these are variants of Power-Of-2, Symmetric, Per-Tensor Quantization, depending on how the parameters are adjusted for Quantization.<br>
-- Simple Calib: Calibration based on min/max ranges
-- Adv Calib: Includes histogram based ranges, calibration of weight/bias parameters to compensate for quantization accuracy loss.
-- Adv DW Calib: Also includes Per-Channel Weight Quantization for Depthwise layers
-- QAT: Quantization Aware Training with PyTorch Jacinto AI DevKit (Does not use Per-Channel Weight Quantization)
+- Simple Calib: Calibration based on min/max ranges<br>
+- Adv Calib: Includes histogram based ranges, calibration of weight/bias parameters to compensate for quantization accuracy loss.<br>
+- Adv DW Calib: Also includes Per-Channel Weight Quantization for Depthwise layers<br>
+- QAT: Quantization Aware Training with PyTorch Jacinto AI DevKit (Does not use Per-Channel Weight Quantization)<br>
 
 **Conclusion based on Simulation Results:**<br>
-- Advanced Calibration Methods may have >2% Accuracy Drop in some cases.
-- Quantization Aware Training (QAT) is consistently able to produce <2% Accuracy drop.
+- Advanced Calibration Methods may have >2% Accuracy Drop in some cases.<br>
+- Quantization Aware Training (QAT) is consistently able to produce <2% Accuracy drop.<br>
 
 
 ## References 
@@ -177,3 +214,5 @@ All of these are variants of Power-Of-2, Symmetric, Per-Tensor Quantization, dep
 
 [8] QUANTIZATION / Introduction to Quantization, https://pytorch.org/docs/stable/quantization.html
 
+[9] Designing Network Design Spaces, Ilija Radosavovic Raj Prateek Kosaraju Ross Girshick Kaiming He Piotr Dollar´, Facebook AI Research (FAIR), https://arxiv.org/pdf/2003.13678.pdf, https://github.com/facebookresearch/pycls
+