Required elements for Quantification Aware Training (QAT)
Description of Quantization aware training (QAT)
There are 2 methods which were implemented in N2D2 and used in many projects and use-cases :
In N2D2 we had Quantizer object which was inserted into initialize/propagate/backpropagate/update methods of Cells which we wanted to quantize.
We quantized weights (of Conv and FC layers) and activations (the easiest one to quantized with these methods is ReLu). The example of the QAT logic is shown in 2 figures below for Conv-BN-ReLu layers.
The propagate method was modified as :
And the backpropagate :
The calculations to transfrom weights/activations into quantized ones, as well as their gradients, for SAT and for LSQ, are CUDA kernels in N2D2. Please refer to this code, as it is the most tested and validated one.
Required elements before starting QAT
Learning process on GPU has to be validated for Conv, BN and FC Cells, for example by training from scratch MobileNet-v1 on ImageNet and getting the SoTa result in Aidge.
Elements needed for the task:
- ImageNet support
- Weight intitialization methods
- Propagate/backpropagate methods for Conv, BN, FC cells in GPU tested and validated. The make sure your backward works don't forget to create checkGradient method, for example for conv
- Loss for training : x-entropy, for example
- Update method validated for Conv, BN, FC cells in GPU using the solver (SGD)
Some useful information about developing QAT :
-
You can do unit tests to check the formulas which you put in CUDA kernel (SAT test, and for LSQ test). It can be a first step to validate kernels, but it will not show all possible issues you might have in the full quantization aware training procedure. The only way to know is to quantize fully your network, for example MobileNet-v1 to 4 bits and check the results against N2D2 ones.
-
There is also pytorch version of SAT and LSQ avalible, if you want to check it out