文章轉(zhuǎn)載于微信公眾號(hào):GiantPandaCV
作者: Pui_Yeung

前言

??量化感知訓(xùn)練（Quantization Aware Training ）是在模型中插入偽量化模塊（fake/_quant module）模擬量化模型在推理過程中進(jìn)行的舍入（rounding）和鉗位（clamping）操作，從而在訓(xùn)練過程中提高模型對(duì)量化效應(yīng)的適應(yīng)能力，獲得更高的量化模型精度 。在這個(gè)過程中，所有計(jì)算（包括模型正反向傳播計(jì)算和偽量化節(jié)點(diǎn)計(jì)算）都是以浮點(diǎn)計(jì)算實(shí)現(xiàn)的，在訓(xùn)練完成后才量化為真正的int8模型。??

Pytorch官方從1.3版本開始提供量化感知訓(xùn)練API，只需修改少量代碼即可實(shí)現(xiàn)量化感知訓(xùn)練。目前torch.quantization仍處于beta階段，不保證API前向、后向兼容性。以下介紹基于Pytorch 1.7，其他版本可能會(huì)有差異。

Pytorch量化感知訓(xùn)練流程

??首先給出提供一個(gè)可運(yùn)行demo，直觀了解量化感知訓(xùn)練的6個(gè)步驟，再進(jìn)行詳細(xì)的介紹

importtorchfromtorch.quantizationimportprepare_qat,get_default_qat_qconfig,convertfromtorchvision.modelsimportquantization# Step1：修改模型#這里直接使用官方修改好的MobileNetV2，下文會(huì)對(duì)修改點(diǎn)進(jìn)行介紹model=quantization.mobilenet_v2()print("originalmodel:")print(model)# Step2：折疊算子#fuse_model()在training或evaluate模式下算子折疊結(jié)果不同，#對(duì)于QAT，需確保在training狀態(tài)下進(jìn)行算子折疊assertmodel.trainingmodel.fuse_model()print("fusedmodel:")print(model)#Step3:指定量化方案#通過給模型實(shí)例增加一個(gè)名為"qconfig"的成員變量實(shí)現(xiàn)量化方案的指定#backend目前支持fbgemm和qnnpackBACKEND="fbgemm"model.qconfig=get_default_qat_qconfig(BACKEND)# Step4：插入偽量化模塊prepare_qat(model,inplace=True)print("modelwithobservers:")print(model)#正常的模型訓(xùn)練，無需修改代碼# Step5：實(shí)施量化model.eval()#執(zhí)行convert函數(shù)前，需確保模型在evaluate模式model_int8=convert(model)print("quantizedmodel:")print(model_int8)# Step6：int8模型推理#指定與qconfig相同的backend，在推理時(shí)使用正確的算子torch.backends.quantized.engine=BACKEND#目前Pytorch的int8算子只支持CPU推理,需確保輸入和模型都在CPU側(cè)#輸入輸出仍為浮點(diǎn)數(shù)fp32_input=torch.randn(1,3,224,224)y=model_int8(fp32_input)print("output:")print(y)

Step1：修改模型

??Pytorch下需要適當(dāng)修改模型才能進(jìn)行量化感知訓(xùn)練，以下以常用的MobileNetV2為例。官方已修改好的MobileNetV2的代碼，詳見這里（https://github.com/pytorch/vi...）

修改主要包括3點(diǎn)，以下摘取相應(yīng)的代碼進(jìn)行介紹：

（1）在模型輸入前加入QuantStub()，在模型輸出后加入DeQuantStub()。目的是將輸入從fp32量化為int8，將輸出從int8反量化為fp32。模型的/_/_init/_/_()和forward()修改為：

classQuantizableMobileNetV2(MobileNetV2):def__init__(self,*args,**kwargs):"""MobileNetV2mainclassArgs:InheritsargsfromfloatingpointMobileNetV2"""super(QuantizableMobileNetV2,self).__init__(*args,**kwargs)self.quant=QuantStub()self.dequant=DeQuantStub()defforward(self,x):x=self.quant(x)x=self._forward_impl(x)x=self.dequant(x)returnx

（2）對(duì)加法等操作加入偽量化節(jié)點(diǎn)。因?yàn)閕nt8數(shù)值進(jìn)行加法運(yùn)算容易超出數(shù)值范圍，所以不是直接進(jìn)行計(jì)算，而是進(jìn)行反量化->計(jì)算->量化的操作。以InvertedResidual的修改為例：

classQuantizableInvertedResidual(InvertedResidual):def__init__(self,*args,**kwargs):super(QuantizableInvertedResidual,self).__init__(*args,**kwargs)#加法的偽量化節(jié)點(diǎn)需要記錄所經(jīng)過該節(jié)點(diǎn)的數(shù)值的范圍，因此需要實(shí)例化一個(gè)對(duì)象self.skip_add=nn.quantized.FloatFunctional()defforward(self,x):ifself.use_res_connect:#普通版本MobileNetV2的加法#returnx+self.conv(x)#量化版本MobileNetV2的加法returnself.skip_add.add(x,self.conv(x))else:returnself.conv(x)

（3）將ReLU6替換為ReLU。MobileNet V2使用ReLU6的原因是對(duì)ReLU的輸出范圍進(jìn)行截?cái)嘁跃徑饬炕癁閒p16模型時(shí)的精度下降。因?yàn)閕nt8量化本身就能確定截?cái)嚅撝?，所以將ReLU6替換為ReLU以去掉截?cái)嚅撝倒潭?的限制。官方的修改代碼在建立網(wǎng)絡(luò)后通過/_replace/_relu()將MobileNetV2中的ReLU6替換為ReLU：

model=QuantizableMobileNetV2(block=QuantizableInvertedResidual,**kwargs)_replace_relu(model)

Step2：算子折疊

??算子折疊是將模型的多個(gè)層合并成一個(gè)層，一般用來減少計(jì)算量和加速推理。對(duì)于量化感知訓(xùn)練而言，算子折疊作用是將模型變“薄”，減少中間計(jì)算過程的誤差積累。

??以下比較有無算子折疊的結(jié)果（上：無算子折疊，下：有算子折疊，打印執(zhí)行prepare/_qat()后的模型）

?如果不進(jìn)行算子折疊，每個(gè)Conv-BN-ReLU單元一共會(huì)插入4個(gè)FakeQuantize模塊。而進(jìn)行算子折疊后，原來Conv2d()被ConvBnReLU2d()代替（3層合并到了第1層），BatchNorm2d()和ReLU()被Inentity()代替（僅作為占位），最終只插入了2個(gè)FakeQuantize模塊。

FakeQuantize模塊的減少意味著推理過程中進(jìn)行的量化-反量化的次數(shù)減少，有利于減少量化帶來的性能損失。

??算子折疊由實(shí)現(xiàn)torch.quantization.fuse/_modules()。目前存在的比較遺憾的2點(diǎn)：

??算子折疊不能自動(dòng)完成，只能由程序員手工指定要折疊的子模型。以torchvision.models.quantization.mobilenet/_v2()中實(shí)現(xiàn)的算子折疊函數(shù)為例：

deffuse_model(self):#遍歷模型內(nèi)的每個(gè)子模型，判斷類型并進(jìn)行相應(yīng)的算子折疊forminself.modules():iftype(m)==ConvBNReLU:fuse_modules(m,['0','1','2'],inplace=True)iftype(m)==QuantizableInvertedResidual:#調(diào)用子模塊實(shí)現(xiàn)的fuse_model()，間接調(diào)用fuse_modules()m.fuse_model()

??能折疊的算子組合有限。目前支持的算子組合為：ConV + BN、ConV + BN + ReLU、Conv + ReLU、Linear + ReLU、BN + ReLU。如果嘗試折疊ConvTranspose2d、ReLU6等不支持的算子則會(huì)報(bào)錯(cuò)。

Step3：指定量化方案

??目前支持fbgemm和qnnpack兩種backend方案。官方推薦x86平臺(tái)使用fbgemm方案，ARM平臺(tái)使用qnnpack方案。??量化方案通過如下方法指定

model.qconfig=get_default_qat_qconfig(backen)#或model.qconfig=get_default_qat_qconfig(backen)

??即通過給model增加一個(gè)名為qconfig為成員變量并賦值。

??量化方案可通過設(shè)置qconfig自定義，本文暫不討論。

Step4：插入偽量化模塊??

通過執(zhí)行prepare/_qat()，實(shí)現(xiàn)按qconfig的配置方案給每個(gè)層增加FakeQuantize()模塊?每個(gè)FakeQuantize()模塊內(nèi)包含相應(yīng)的Observer()模塊，在模型執(zhí)行forward()時(shí)自動(dòng)記錄數(shù)值，供實(shí)施量化時(shí)使用。

Step5：實(shí)施量化??

完成訓(xùn)練后，通過執(zhí)行convert()轉(zhuǎn)換為真正的int8量化模型。?完成轉(zhuǎn)換后，F(xiàn)akeQuantize()模塊被去掉，原來的ConvBNReLU2d()算子被替換為QuantizedConvReLU2d()算子。

Step6：int8模型推理

??int8模型的調(diào)用方法與普通的fp32模型的調(diào)用無異。需要注意的是，目前量化算子僅支持CPU計(jì)算，故須確保輸入和模型都在CPU側(cè)。

??若模型推理中出現(xiàn)報(bào)錯(cuò)，一般是前面的步驟存在設(shè)置不當(dāng)，參考常見問題第1點(diǎn)。

常見問題

(1) RuntimeError: Could not run XX with arguments from the YY backend. XX is only available for these backends ZZ??

雖然fp32模型和int8模型都能在CPU上推理，但fp32算子僅接受tensor作為輸入，int8算子僅接受quantedtensor作為輸入，輸入和算子的類型不一致導(dǎo)致上述錯(cuò)誤。

??一般排查方向?yàn)椋菏欠裢瓿闪四Ｐ托薷?，將加法等操作替換為量化版本；是否正確添加了QuantStub()和DeQuantStub()；是否在執(zhí)行convert()前是否執(zhí)行了model.eval()（在traning模型下，dropout無int8實(shí)現(xiàn)但沒有被去掉，然而在執(zhí)行推理時(shí)會(huì)報(bào)錯(cuò)）。

(2) 是否支持GPU訓(xùn)練，是否支持DistributedDataParallel訓(xùn)練？??

支持。官方有一個(gè)完整的量化感知訓(xùn)練的實(shí)現(xiàn)，使用了GPU和DistributedDataParallel，可惜在文檔和教程中未提及，參考這里（https://github.com/pytorch/vi.../_quantization.py）。

(3) 是否支持混合精度模型（例如一部分fp32推理，一部分int8推理）？??

官方?jīng)]有明確說明，但經(jīng)實(shí)踐是可以的。

??模型是否進(jìn)行量化取決于是否帶qconfig。因此可以將模型定義修改為

classMixModel(nn.Module):def__init__(self):super(MixModel,self).__init__()self.fp32_part=Fp32Model()self.int8_part=Int8Model()defforward(self,x):x=self.int8_part(x)x=self.fp32(x)returnxmix_model=MixModel()mix_model.int8_part.qconfig=get_default_qat_qconfig(BACKEND)prepare_qat(mix_model,inplace=True)

??由此可實(shí)現(xiàn)所需的功能。注意將QuantStub()、Dequant()模塊移到Int8Model()中。

（4）精度保持效果如何，如何提升精度？??

筆者進(jìn)行的實(shí)驗(yàn)不多，在做過的簡(jiǎn)單的OCR任務(wù)中，可以做到文字檢測(cè)和識(shí)別模型的指標(biāo)下降均不超過1個(gè)點(diǎn)（量化的int8模型對(duì)比正常訓(xùn)練的fp32模型）。官方教程中提供了分類例子的效果和提升精度的技巧，可供參考。

總結(jié)

??Pytorch官方提供的量化感知訓(xùn)練API，上手較為簡(jiǎn)單，易于集成到現(xiàn)有訓(xùn)練代碼中。但目前手動(dòng)修改模型和算子折疊增加了一定的工作量，期待在未來版本的改進(jìn)。

- END -

推薦閱讀

• PyTorch OCR模型的安卓端部署
• 深度學(xué)習(xí)量化技術(shù)科普
• 簡(jiǎn)單粗暴的多對(duì)象目標(biāo)跟蹤神器 – DeepSort

更多嵌入式AI技術(shù)干貨請(qǐng)關(guān)注嵌入式AI專欄。

審核編輯：符乾江