HCHUNGW的部落格

12.8 QAT 量化實踐

前言

上一小節介紹了PTQ量化，本節介紹另外一種重要量化方法——QAT（quantization aware training，感知量化訓練）。

本節將從QDQ節點出發，深入瞭解QAT的量化過程，然後基於QAT進行ResNet50的量化及性能、效率評估。

QAT基礎概念

再回顧上上小結QAT與PTQ的區別。

• PTQ（Post-training quantization，訓練後量化），是不需要訓練資料（帶標籤的資料），通常採用小部分資料（不需要標籤）進行模型啟動值校準，統計啟動值的分佈，然後計算動態範圍，再計算量化scale和Z值；對於weight和bias直接根據模型權重資料進行量化。
• QAT(Quantization-aware training，量化感知訓練) ，是需要訓練資料，並需要在訓練階段插入偽量化層，訓練得到帶QDQ節點的模型，然後在量化時採用QDQ節點進行量化，通常可以得到比PTQ更高的精度。
•  
  • QDQ（quantize節點、dequantize節點）節點是QAT的靈魂，是一個量化節點和反量化節點，可在訓練時進行反覆運算優化，思想是在訓練的時候獲得一個好的Q節點，它的scale和Z值是比較好的scale和Z值，為什麼說它得到的值比較好呢？因為這些值可通過DQ節點恢復出更接近真實值的資料，因此認為訓練階段獲得的Q節點中的scale和Z值是較好的選擇。
  • 

偽量化節點

剛接觸量化時，偽量化是較難理解的，偽量化的思想是在訓練時模擬量化帶來的誤差，通過反向傳播來對抗誤差帶來的精度下降，從而得到一個較好的scale和Z值。

具體可參考Nvidia的PPT《toward-int8-inference-deploying-quantization-aware-trained-networks-using-tensorrt》

下圖是一個常規fp32的運算過程，接下來將逐步講解QAT過程。

第一步，插入QDQ（fake quant ops）節點。

插入QDQ的地方即需要量化的地方，回顧上上小結，需要量化的通常是weight、bias和activation，因此有以下4個節點需要量化：

X：可看成第一個activation，因此是需要量化的

Wx：conv1的權重

Relu：之所以在Relu而不是conv1之後，是因為Relu是線性的，可以合併它兩，在Relu後做FQ即可。

Wy：conv2的權重

第二步，訓練匯出ONNX，最終得到scale和Z值。

第三步，在TensorRT中對計算圖進行優化

• 常量的折疊：如權重的Q節點可與權重合並，無需在真實推理中由fp32的權重經過scale和Z變為int8的權重

• op融合：將DQ資訊融合到運算元（如圖中conv）中，通過op融合，模型計算將變為真實的int8輸入、int8輸出

于TensorRT是如何做op融合的，暫時找不到相關資料，暫且略過。

思想大概能理解，QConvRelu是吸收了DQ節點的資訊，改變weight的數值，這樣能更好的計算int8與int8的卷積。

通過以上案例可知conv2輸出的是fp32，因為它沒有插入QDQ，而conv1輸出的是int8，因為插入了QDQ。

這點可以讓我們在量化時根據需要，自訂插入QDQ，使得需要高精度輸出的地方保持fp32。大意如下圖所示

總結一下，QAT的過程：

1. 插入QDQ
2. 訓練，得到scale和Z值，匯出onnx模型
3. TensorRT圖優化、運算元融合，匯出engine，運算元融合中將DQ資訊注入op中，使得op接收int8，輸出int8，從而提升了運算速度。

ResNet50 QAT 實踐

接下來進行QAT實踐，初識量化時，一直認為QAT比PTQ繁瑣，因為QAT需要訓練，而PTQ直接校準就好了。但是，在代碼實現上，QAT反而比PTQ簡單許多。

流程分析

首先，分析QAT的實現流程步驟：

• 第一步，初始化pytorch_quantization，將pytorch的模組替換掉，這一步pq庫已經封裝好了，只需要一行代碼quant_modules.initialize()
• 第二步，創建模型、載入訓練權重
• 第三步，常規訓練步驟
• 第四步，保存成ONNX模型和pth模型

由此可知，只是在常規訓練流程的代碼中，插入quant_modules.initialize() 即可。

當然，在QAT訓練和常規訓練不一樣，學習率通常要小100倍，同時可以用余弦下降法，逐步下降學習率。

TensorRT文檔中給了一些QAT的建議：

Usually, it doesn’t need to fine tune very long. We usually use around **10% of the original training schedule**, starting at **1% of the initial training learning rate**, and a cosine annealing learning rate schedule that follows the decreasing half of a cosine period, down to 1% of the initial fine tuning learning rate (0.01% of the initial training learning rate).

Copy

代碼說明

同樣的，代碼位於chapter-8/01_classification中，代碼整體與訓練代碼train_main.py保持一致，其中刪減了一些不必要的日誌記錄程式碼片段。

同時修改了學習率調整方法，以及增加ONNX模型匯出。

這裡採用4epoch，lr=0.001*0.01進行訓練，代碼直接運行即可：python resnet50_qat.py

日誌如下：

Epoch: [000/005]  Train Loss avg: 0.1086  Valid Loss avg: 0.1770  Train Acc@1 avg:  96.1009   Valid Acc@1 avg: 93.4295

Epoch: [001/005]  Train Loss avg: 0.0786  Valid Loss avg: 0.1819  Train Acc@1 avg:  97.2477   Valid Acc@1 avg: 93.4295

Epoch: [002/005]  Train Loss avg: 0.0773  Valid Loss avg: 0.1742  Train Acc@1 avg:  97.3624   Valid Acc@1 avg: 92.9487

Epoch: [003/005]  Train Loss avg: 0.0735  Valid Loss avg: 0.1771  Train Acc@1 avg:  97.0183   Valid Acc@1 avg: 93.5897

Epoch: [004/005]  Train Loss avg: 0.0633  Valid Loss avg: 0.1704  Train Acc@1 avg:  97.9740   Valid Acc@1 avg: 93.4295

Copy

可以看到，Accuracy基本穩定在93.5上下，最優為93.58，相較於原始的94.3，掉了不到1個百分點，比PTQ穩定以及性能更優。

效率對比

得到onnx模型，下面比較fp32、PTQ量化和QAT量化，三者的時延和輸送量。

可以看到PTQ和QAT在效率上是一樣的。

時延(中位數) ms：

輸送量(FPS)

FP32\ptq\qat 三種模型的對比指令如下：

trtexec --onnx=resnet_50_fp32_bs1.onnx

trtexec --onnx=resnet_50_fp32_bs32.onnx

trtexec --onnx=resnet_50_ptq_bs1_data-num512_percentile_91.03%.onnx --int8

trtexec --onnx=resnet_50_ptq_bs32_data-num512_percentile_91.03%.onnx --int8 --saveEngine=resnet_ptq_int8.engine

trtexec --onnx=resnet_50_qat_bs1_93.43%.onnx --int8

trtexec --onnx=resnet_50_qat_bs32_93.43%.onnx --int8

Copy

注1：在使用trtexec時，看到了這樣的日誌，隨後嘗試了加上 --fp16，整體輸送量又提升了~40%（相較於fp32提高140%），這或許是將fp32變為fp16帶來的效率提升？

[09/29/2023-16:20:37] [I] FP32 and INT8 precisions have been specified - more performance might be enabled by additionally specifying --fp16 or --best

trtexec --onnx=resnet_50_qat_bs32_93.43%.onnx --int8 --best

trtexec --onnx=resnet_50_qat_bs32_93.43%.onnx --int8 --fp16

Copy

注2：在新電腦裡又看到了這樣的warning，嘗試使用--useCudaGraph ，可以提高吞吐。嘗試之後，發現吞吐從1200 提升到了1600，看來trtexec裡的提示資訊非常有説明！ [04/04/2024-20:55:04] [W] Throughput may be bound by Enqueue Time rather than GPU Compute and the GPU may be under-utilized. [04/04/2024-20:55:04] [W] If not already in use, --useCudaGraph (utilize CUDA graphs where possible) may increase the throughput.*

小結

本節回顧了QAT和PTQ的區別，同時介紹了QDQ（偽量化運算元：量化和反量化）在訓練、推理時的變化，最後通過代碼實現QAT，並對比了性能和效率變化。

到這裡，量化的基礎就結束了，模型量化的內容遠不止這些，模型量化還有更高級的內容，例如逐層精度損失分析、逐層插入/取消量化節點、自訂量化演算法。

想要深入研究模型加速/量化的朋友，可以進一步學習其它資料,這裡推薦一些學習量化用到的資料：

1. TRT int8量化文檔：https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#working-with-int8
2. TRT範例代碼：https://github.com/NVIDIA/TensorRT/blob/main/quickstart/quantization_tutorial/qat-ptq-workflow.ipynb
3. pytorch-quantization：https://github.com/NVIDIA/TensorRT/tree/main/tools/pytorch-quantization
4. nvidia int8 QAT量化介紹：https://developer.nvidia.com/blog/achieving-fp32-accuracy-for-int8-inference-using-quantization-aware-training-with-tensorrt/
5. yolov6 量化工程範例：https://tech.meituan.com/2022/09/22/yolov6-quantization-in-meituan.html
6. 老潘的TRTv8量化入門博客：https://zhuanlan.zhihu.com/p/479101029
7. B站量化系列教程-ZOMI醬：https://www.bilibili.com/video/BV1VD4y1n7AR/?spm_id_from=333.788&vd_source=19b783e279d4d4ceff5b927f27ea6aa3
8. B站量化系列教程-手寫AI：https://www.bilibili.com/video/BV18L41197Uz/?spm_id_from=333.337.search-card.all.click&vd_source=19b783e279d4d4ceff5b927f27ea6aa3
9. MIT的TinyML-第六節：MIT-TinyML-Lec06-Quantization-II

下一小節，將利用本章學習的TensorRT知識，進行YOLOv8的量化實踐，學習優秀開原始程式碼的工程化過程。