量化感覺訓練實踐：實作精度無損的模型壓縮和推理加速

1. 概述

對深度學習模型進行低比特量化，可以有效地降低模型部署時在存儲、計算、通信上的開銷，是一種常見的模型壓縮和推理優化技術。然而，模型量化在實際應用中仍然存在不少挑戰，最為常見的問題就是模型精度下降（如無特殊說明，本文中“模型精度”是指準确率等模型應用于具體任務的效果名額）。以計算機視覺領域為例，在目标檢測、圖像分割等複雜任務上，量化帶來的精度下降更為明顯。

通過在模型訓練階段引入量化相關限制，即量化感覺訓練（Quantization-aware training，QAT），能夠更好地解決模型量化的精度問題。本文以近期流行的YOLOX[8]目标檢測模型為例，介紹量化感覺訓練的原理流程，讨論如何實作精度無損的實踐經驗，并展示了量化後的模型能夠做到精度不低于原始浮點模型，模型壓縮4X、推理加速最高2.3X的優化效果。

2. 量化原理

在數字信号處理領域，量化是指将信号的連續取值（或者大量可能的離散取值）近似為有限多個（或較少的）離散值的過程。具體到深度學習領域，模型量化是指将浮點激活值或權重（通常以32比特浮點數表示）近似為低比特的整數（16比特或8比特），進而在低比特的表示下完成計算的過程。通常而言，模型量化可以壓縮模型參數，進而降低模型存儲開銷；并且通過降低訪存和有效利用低比特計算指令等，能夠取得推理速度的提升，這對于在資源受限裝置上部署模型尤為重要。

給定浮點類型的值，可以通過如下公式将它轉化成8比特量化值：

其中，

表示量化的scale，

與

分别表示量化值域的最小值與最大值，

表示輸入浮點值，

表示量化後的值。量化值轉化為浮點值隻需執行反操作即可：

進一步的，在将輸入資料和權重進行量化後，我們就能夠将神經網絡的常見操作轉換為量化操作。以卷積操作為例，其量化版本典型的計算流程如圖1所示：

圖1 典型的量化卷積算子計算流程圖

1. 權重與輸入先量化成8bit，進行卷積操作，用32bit

   

   來存儲中間結果；
2. bias量化為32bit，與

   

   進行相加為

   

   ；
3. 利用

   

   、

   

   以及

   

   的量化scale将32bit的

   

   轉化為8bit的

   

4. 如果該層的下一層也是量化OP，則

   

   可直接輸出給下一層；如果是非量化OP，則将

   

   反量化為浮點值後，再輸出給下一層。

從上述量化計算的原理能夠容易看出，将浮點數轉化為低比特整數進行計算會不可避免地引入誤差，神經網絡模型中每層量化計算的誤差會累積為模型整體精度的誤差。常見的訓練後量化（Post Training Quantization，PTQ）方案中，通過統計在典型輸入資料情況下，待量化變量的數值分布，來選擇合适的量化參數（scale，zero point等），将因量化而引入的資訊損失降低到最小。

但是PTQ方案往往還是無法實作精度無損的模型量化，為了進一步降低量化帶來的精度下降，我們可以采用量化感覺訓練的方案，在訓練的計算圖中引入僞量化的操作，通過微調訓練(finetuning)讓模型權重“适應”量化引入的誤差，以實作更好的、甚至無損的量化模型精度。

3. YOLOX量化訓練

我們以YOLOX-s目标檢測模型(GitHub repo[1])為例，使用公開的預訓練模型參數，在COCO2017資料集上進行量化訓練實驗。量化訓練算法選擇LSQ[2,3]，該系列算法利用梯度來更新量化的scale與zero_point，在不需要精細調節參數的情況下能夠獲得較好的性能。為了通過量化訓練獲得更好的量化模型精度，我們需要重點關注如下幾點設定：

3.1 與部署後端相比對的量化方式

不同的部署後端，可能采用不用的量化計算實作方式，需要比對訓練和部署階段的量化方式以避免引入額外誤差。以PyTorch[7]預設的CPU後端為例，基本的量化方式為

• • weight： per-channel，int8，對稱量化
  • activation： per-tensor，uint8，非對稱量化
  • 量化對象： 所有的Conv2d

以移動端架構MNN為例，基本的量化方式為：

• • activation： per-tensor，int8，對稱量化

通常在具體推理架構上部署量化模型時，還會對類似

conv-bn

conv-relu/relu6

conv-bn-relu/relu6

這樣的卷積層進行算子融合，是以需要設定為量化融合後的activation。YOLOX中的SiLU激活函數通常不會被融合，output activation的量化僅需要設定到bn的輸出，而不是SiLU的輸出。訓練時一個典型的卷積層量化位置如圖2所示。

圖2 量化位置示意圖

同時，QAT時所有的BN全部fold入相應的卷積中，實驗采取了文獻[6]中的fold政策。由于模拟量化的引入可能會使得BN層的running_mean與running_var不穩定，進而導緻量化訓練無法收斂。是以，我們在量化訓練中固定BN的running_mean與running_var。

此外，特定部署後端的實作可能需要将activation量化為7bit數值範圍，以防止計算溢出。帶有avx512_vnni指令集的CPU型号上，則沒有相應要求。

3.2 量化訓練參數初始化

為了避免導緻訓練loss不穩定甚至訓練發散的情況，我們利用訓練後量化（Post training quantization，PTQ）所得的量化參數來初始化LSQ量化訓練中activation scale參數。基本的步驟是：

• • 選取幾十張典型圖檔，在預訓練模型上進行推理，并收集統計各個待量化的activation資訊。
  • 使用如MSE、KL散度等metric來計算各activation的最佳scale。

基本的實作方式可以使用PyTorch的

HistogramObserver

計算activation scale & zero point，PerChannelMinMaxObserver計算weight scale。

3.3 訓練超參數

在QAT階段，我們使用已經收斂的預訓練模型權重，并且使用PTQ量化參數進行QAT初始化。這種情況下，模型參數已接近收斂，是以我們将整體訓練的超參數和原始YOLOX訓練保持一緻，将學習率設定為原始訓練收斂階段的學習率，即5e-4。

3.4 特定後端算子實作的計算誤差

因為數值表示精度問題，round操作的結果在訓練架構與推理後端上可能并不是相同的，例如：

import torch

torch.tensor(2.5).cuda().round()  # 輸出tensor(2., device='cuda:0')

torch.tensor(3.5).cuda().round()  # 輸出tensor(4., device='cuda:0')

2.5和3.5的四舍五入行為在PyTorch上并不是相同的，而該行為差異在常見的部署後端架構上可能不存在。這種問題會導緻QAT模拟精度和後端實際運作時的精度存在明顯差異，是以需要在量化訓練階段進行修正。例如針對MNN後端，我們可以采取如下操作來避免這個差異。

def round_pass(x):

    """

    A simple way to achieve STE operation.

    # y = torch.round(x) # for PyTorch backend

    y = (x + torch.sign(x) * 1e-6).round()  # for mnn backend, round前添加一個小的數值

    y_grad = x

    return (y - y_grad).detach() + y_grad

4. 實驗結果與分析

4.1 精度和加速效果

按照上述方式在YOLOX-s模型上進行量化訓練，使用COCO2017 validation set進行精度驗證，結果如表1所示。兩種後端上真實量化的模型的精度性能均和浮點模型齊平。

量化方式	模型 （YOLOX-s）	mAP @0.5:0.95	@0.5	@0.75
-	浮點模型	40.5	59.3	43.8
權重：對稱量化 激活值：對稱量化	PTQ	39.3	58.5	43.0
QAT	40.8	59.9	44.1
真實量化模型	40.6	59.7	44.2
激活值：非對稱量化		39.9	58.9	43.2
	40.7	59.8	43.9
			43.7

表1 浮點模型、QAT模型及後端真實量化模型的精度對比

速度實驗中，我們選取PyTorch[7]後端及x86測試平台進行測試，測試的圖檔分辨率為1x3x640x640。不同資源數下的結果如表2所示。在精度無損的前提下，量化模型的推理速度最高可以提升2.35x，同時模型尺寸為原來的1/4。

後端	浮點速度 (ms)	量化速度	加速比	裝置	線程數
PyTorch	321.2	189.9	1.70x	Intel(R) Xeon(R) Platinum 8369HC CPU @ 3.30GHz	1
218.3	106.2	2.06x	2
117.3	57.7	2.03x	4
75.3	34.1	2.21x	8
57.9	24.7	2.35x	16

表2 浮點模型、QAT模型及後端真實量化模型的速度對比

4.2 量化參數初始化的影響

LSQ及LSQ+論文中提出了相應的量化資訊初始化方法，實際使用中會發現該種初始化方法對于學習率的設定會比較敏感。一方面，若量化參數的初始值距離收斂值較遠，需要通過設定較大的學習率來訓練他們。另一方面，較大的學習率會使得模型本身參數跳出已收斂的狀态，使得模型進入“重訓練”的狀态，不确定性增加。而用PTQ來初始化量化參數，可以使得參數處于較優的初始狀态，利用浮點模型收斂時的學習率去進行finetune即可獲得較好的性能。固定學習率為4e-5的前提下，PTQ及LSQ初始化的訓練結果與曲線如圖4所示。LSQ初始化的訓練曲線在開啟mossac資料增強時是逐漸向下的，而PTQ初始化是逐漸向上的。

圖4 固定finetune學習率下，不同初始化方法的訓練曲線

4.3 訓練超參數的影響

學習率的設定會直接影響到最終QAT模型的性能。以原始模型訓練收斂階段的學習率（5e-4）為基準，如果QAT階段使用相同的學習率，QAT初期的模型精度會逐漸下降，如圖4中PTQ初始化紅色曲線訓練早期所示，但是最終精度會提升至原始非量化模型的水準。如果QAT使用更低學習率（例如5e-4），模型精度會相比于PTQ初始化狀态逐漸上升，但是最終精度提升不大：

學習率	mAP@50:95
Baseline
		39.4
	5e-6
	7.5e-6	39.8
	2.5e-6

上述現象的一種可能原因是，在國小習率下模型權重、量化scale基本不變。實際上是基于PTQ的初始解空間繼續向局部更好的收斂點靠近，跳出局部解空間的可能性較低。是以，在PTQ初始化的前提下，學習率可以直接設定成浮點模型訓練收斂階段的值。

4.4 訓練輪數的選擇

上述QAT的結果是訓練300 epochs後的模型，減少訓練epoch數量的結果，如下所示：

	AP @50:95
QAT（300 epoch）	5e-4
QAT（15 epoch）		39.6		43.3
QAT（30 epoch）			58.8

可以看出随着訓練輪數的增大，QAT的結果會更好。QAT訓練輪數較低，結果會不如直接用國小習率進行finetune。我們可以得出經驗性的trade-off：如果計算資源充足，可以選擇訓練更長的時間來獲得更好的性能。如果計算資源較少，則可以用國小習率訓練較短的時間，以獲得高于PTQ的性能。

4.5 修正特定算子計算誤差的影響

特定算子在訓練架構與後端架構中的行為會有細微的差别，這會導緻量化訓練精度與實際量化模型的精度産生較大的差異。以MNN[5]為例，在YOLOX模型中存在如下兩類OP會導緻該現象。

修正round操作的影響

是否對round進行修正的結果如表3所示。從訓練角度而言，round經過修正後的性能會略好于未修正的。同時如果在訓練時round的行為與後端不一緻的話，可能會導緻真實量化模型的精度發生較大的變化。

修正
未修正			32.9

表3 round修正對于量化訓練的影響

Sigmoid快速實作引入的誤差

為了提升指數計算的速度，後端架構通常會采取一些快速近似計算。這對于浮點模型而言，通常不會引入較大的誤差。但是對于量化模型而言，這個誤差可能因為量化操作而被放大。

如圖3所示，對于YOLOX的主要pattern（Conv -> SiLU -> Conv），前一層Conv的輸出經過SiLU函數後，快速近似計算引入的誤差會被後一層卷積輸入處的量化操作（除以scale）而縮放。scale越小，縮放的程度越大。以YOLOX-s為例，是否對指數計算進行近似的量化模型精度如表4所示。

圖3 YOLOX量化pattern

近似	40.2	59.5
未近似

表4 指數計算的近似帶來的模型性能誤差

5. 總結

本文通過在YOLOX目标檢測模型上的量化實踐，驗證了通過量化感覺訓練(QAT)能夠在精度無損的情況下，獲得顯著的模型壓縮和推理加速。我們對量化精度誤差因素進行了具體分析，指出了解決精度問題的一系列實踐手段，并通過實驗驗證了效果，這可以作為我們在實際應用模型量化時的經驗參考。雖然量化相關方法已經被大量研究，但是在實際複雜任務中應用量化仍然面臨不少挑戰，真正将量化壓縮落地，更需要通過模型和系統兩方面的協同。此外，還有更多量化訓練相關的技術方案（如混合精度量化、更低比特量化等）值得探索和完善。

關于我們

本文中關于量化訓練實踐工作由阿裡雲-PAI模型壓縮團隊和微軟NNI團隊合作完成，也感謝MNN團隊的技術支援。更多模型壓縮的算法實作可參考NNI（GitHub repo[4]），更多模型推理優化技術方案可見阿裡雲

PAI-Blade

。

參考文獻&代碼倉庫

[1]

https://github.com/Megvii-BaseDetection/YOLOX

[2] Esser S K, McKinstry J L, Bablani D, et al. Learned step size quantization[J]. arXiv preprint arXiv:1902.08153, 2019.

[3] Bhalgat Y, Lee J, Nagel M, et al. Lsq+: Improving low-bit quantization through learnable offsets and better initialization[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops. 2020: 696-697.

[4]

https://github.com/microsoft/nni

[5] Jiang X, Wang H, Chen Y, et al. Mnn: A universal and efficient inference engine[J]. arXiv preprint arXiv:2002.12418, 2020.

[6] Jacob B, Kligys S, Chen B, et al. Quantization and training of neural networks for efficient integer-arithmetic-only inference[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2018: 2704-2713.

[7] Paszke A, Gross S, Massa F, et al. Pytorch: An imperative style, high-performance deep learning library[J]. Advances in neural information processing systems, 2019, 32: 8026-8037.

[8] Ge Z, Liu S, Wang F, et al. Yolox: Exceeding yolo series in 2021[J]. arXiv preprint arXiv:2107.08430, 2021.

若有收獲，就點個贊吧