Pytorch量化入門之超分量化（一）
?點(diǎn)擊上方“AI算法與圖像處理”，選擇加"星標(biāo)"或“置頂”
重磅干貨，第一時(shí)間送達(dá)
來源:AIWakler
最近Happy在嘗試進(jìn)行圖像超分的INT8量化，發(fā)現(xiàn)：pytorch量化里面的坑真多，遠(yuǎn)不如TensorFlow的量化好用。不過花了點(diǎn)時(shí)間終于還是用pytorch把圖像超分模型完成了量化，以EDSR為例，模型大小73%,推理速度提升40%左右(PC端)，視覺效果幾乎無損，定量指標(biāo)待補(bǔ)充。有感于網(wǎng)絡(luò)上介紹量化的博客一堆，但真正有幫助的較少，所以Happy會(huì)盡量以圖像超分為例提供一個(gè)完整的可復(fù)現(xiàn)的量化示例(分兩章內(nèi)容進(jìn)行)。
背景
量化在不同領(lǐng)域有不同的定義，而在深度學(xué)習(xí)領(lǐng)域，量化有兩個(gè)層面的意義：(1) 存儲(chǔ)量化，即更少的bit來存儲(chǔ)原本需要用浮點(diǎn)數(shù)(一般為FP32)存儲(chǔ)的tensor；(2) 計(jì)算量化，即用更少的bit來完成原本需要基于浮點(diǎn)數(shù)(一般為FP32，F(xiàn)P16現(xiàn)在也是常用的一種)完成的計(jì)算。量化一般有這樣兩點(diǎn)好處：
更小的模型體積，理論上減少為FP32模型的75%左右，從筆者不多的經(jīng)驗(yàn)來看，往往可以減少73%； 更少的內(nèi)存訪問與更快的INT8計(jì)算，從筆者的幾個(gè)簡(jiǎn)單嘗試來看，一般可以加速40%左右，這個(gè)還會(huì)跟平臺(tái)相關(guān)。
對(duì)于量化后模型而言，其部分或者全部tensor(與量化方式、量化op的支持程度有關(guān))將采用INT類型進(jìn)行計(jì)算，而非量化前的浮點(diǎn)類型。量化對(duì)于底層的硬件支持、推理框架等要求還是比較高的，目前X86CPU，ARMCPU，Qualcomm DSP等主流硬件對(duì)量化都提供了支持；而NCNN、MACE、MNN、TFLite、Caffe2、TensorRT等推理框架也都對(duì)量化提供了支持，不過不同框架的支持度還是不太一樣，這個(gè)就不細(xì)說了，感興趣的同學(xué)可以自行百度一下。
筆者主要用Pytorch進(jìn)行研發(fā)，所以花了點(diǎn)精力對(duì)其進(jìn)行了一些研究&嘗試。目前Pytorch已經(jīng)更新到了1.7版本，基本上支持常見的op，可以參考如下：
Activation：ReLU、ReLU6、Hardswish、ELU； Normalization：BatchNorm、LayerNorm、GroupNorm、InstanceNorm； Convolution：Conv1d、Conv2d、Conv3d、ConvTranspose1d、ConvTranspose2d、Linear； Other：Embedding、EmbeddingBag。
目前Pytorch支持的量化有如下三種方式：
Post Training Dynamic Quantization：動(dòng)態(tài)量化，推理過程中的量化，這種量化方式常見諸于NLP領(lǐng)域，在CV領(lǐng)域較少應(yīng)用； Post Training Static Quantization：靜態(tài)量化，訓(xùn)練后靜態(tài)量化，這是CV領(lǐng)域應(yīng)用非常多的一種量化方式； Quantization Aware Training：感知量化，邊訓(xùn)練邊量化，一種比靜態(tài)量化更優(yōu)的量化方式，但量化時(shí)間會(huì)更長(zhǎng)，但精度幾乎無損。
注：筆者主要關(guān)注CV領(lǐng)域，所以本文也將主要介紹靜態(tài)量化與感知量化這種方式。
Tensor量化
要實(shí)現(xiàn)量化，那么就不可避免會(huì)涉及到tensor的量化，一般來說，量化公式可以描述如下：
目前Pytorch中的tensor支持int8/uint8/int32等類型的數(shù)據(jù)，并同時(shí)scale、zero_point、quantization_scheme等量化信息。這里，我們給出一個(gè)tensor量化的簡(jiǎn)單示例：
x?=?torch.rand(3,?3)
print(x)
x?=?torch.quantize_per_tensor(x,?scale=0.2,?zero_point=3,?dtype=torch.quint8)
print(x)
print(x.int_repr())
一個(gè)參考輸出如下所示：
注1：藍(lán)框?yàn)樵嫉母↑c(diǎn)數(shù)據(jù)，紅框?yàn)閠ensor的量化信息，綠框則對(duì)應(yīng)了量化后的INT8數(shù)值。
注2：量化不可避免會(huì)出現(xiàn)精度損失，這個(gè)損失與scale、zero_point有關(guān)。
在量化方面，Tensor一般有兩種量化模式：per tensor與per channel。對(duì)于PerTensor而言，它的所有數(shù)值都按照相同方式進(jìn)行scale和zero_point處理；而對(duì)于PerChannel而言，它有多種不同的scale和zero_point參數(shù)，這種方式的量化精度損失更少。
Post Training Static Quantization
靜態(tài)量化一般有兩種形式：（1） 僅weight量化；(2) weight與activation同時(shí)量化。對(duì)于第一種“僅weight量化”而言，只針對(duì)weight量化可以使得模型參數(shù)所占內(nèi)存顯著減小，但在實(shí)際推理過程中仍需要轉(zhuǎn)換成浮點(diǎn)數(shù)進(jìn)行計(jì)算；而第二種“weight與activation同時(shí)量化”則不僅對(duì)weight進(jìn)行量化，還需要結(jié)合校驗(yàn)數(shù)據(jù)進(jìn)行activation的量化。第一種的量化非常簡(jiǎn)單，這里略過，本文僅針對(duì)第二種方式進(jìn)行介紹。
Pytorch的靜態(tài)量化一把包含五個(gè)步驟：
fuse_model：該步驟用來對(duì)可以融合的op進(jìn)行融合，比如Conv與BN的融合、Conv與ReLU的融合、Conv與BN以及ReLU的融合、Linear與BN的融合、Linear與BN以及ReLU的融合。目前Pytorch已經(jīng)內(nèi)置的融合code：
fuse_modules(model,?modules_to_fuse,?inplace=False,?fuser_func=fuse_known_modules,?fuse_custom_config_dict=None)
在完成融合后，第一個(gè)op將被替換會(huì)融合后的op，而其他op則會(huì)替換為nn.Identity。
qconfig：該步驟用于設(shè)置用于模型量化的方式，它將插入兩個(gè)observer，一個(gè)用于監(jiān)測(cè)activation，一個(gè)用于監(jiān)測(cè)weight?？紤]到推理平臺(tái)的不同，pytorch提供了兩種量化配置：針對(duì)x86平臺(tái)的 fbgemm以及針對(duì)arm平臺(tái)的qnnpack。
不同平臺(tái)的量化配置方式存在些微的區(qū)別，大概如下：
| backend | activation | weight |
|---|---|---|
| Fbgemm | HistogramObserver(reduce_range=True) | PerChannelMINMaxObserver(default_per_channel_weight_observer) |
| Qnnpack | HistogramObserver(reduce_range=False) | PerChannelMINMaxObserver(default_per_channel_weight_observer) |
| Default | MinMaxObserver(default_observer) | MinMaxObserver(default_weight_observer) |
Prepare:該步驟用于給每個(gè)支持量化的模塊插入Observer，用于收集數(shù)據(jù)并進(jìn)行量化數(shù)據(jù)分析。以activation為例，它將根據(jù)所喂入數(shù)據(jù)統(tǒng)計(jì)min_val與max_val，一般觀察幾個(gè)次迭代即可，然后根據(jù)所觀察到數(shù)據(jù)進(jìn)行統(tǒng)計(jì)分析得到scale與zero_point。 Feed Data：為了更好的獲得activation的量化參數(shù)信息，我們需要一個(gè)合適大小的校驗(yàn)數(shù)據(jù)，并將其送入到前述模型中。這個(gè)就比較簡(jiǎn)單了，就按照模型驗(yàn)證方式往里面送數(shù)據(jù)就可以了。 Convert：在完成前述四個(gè)步驟后，接下來就需要將完成量化的模型轉(zhuǎn)換為量化后模型了，這個(gè)就比較簡(jiǎn)單了，通過如下命令即可。
torch.quantization.convert(model,?inplace=True)
該過程本質(zhì)上就是用量化OP替換模型中的費(fèi)量化OP，比如用nnq.Conv2d替換nn.Conv2d, nnq.ConvReLU2d替換nni.ConvReLU2d(注：這是Conv與ReLU的合并)。之前的量化op以及對(duì)應(yīng)的被替換op列表如下：
DEFAULT_STATIC_QUANT_MODULE_MAPPINGS?=?{
????QuantStub:?nnq.Quantize,
????DeQuantStub:?nnq.DeQuantize,
????nn.BatchNorm2d:?nnq.BatchNorm2d,
????nn.BatchNorm3d:?nnq.BatchNorm3d,
????nn.Conv1d:?nnq.Conv1d,
????nn.Conv2d:?nnq.Conv2d,
????nn.Conv3d:?nnq.Conv3d,
????nn.ConvTranspose1d:?nnq.ConvTranspose1d,
????nn.ConvTranspose2d:?nnq.ConvTranspose2d,
????nn.ELU:?nnq.ELU,
????nn.Embedding:?nnq.Embedding,
????nn.EmbeddingBag:?nnq.EmbeddingBag,
????nn.GroupNorm:?nnq.GroupNorm,
????nn.Hardswish:?nnq.Hardswish,
????nn.InstanceNorm1d:?nnq.InstanceNorm1d,
????nn.InstanceNorm2d:?nnq.InstanceNorm2d,
????nn.InstanceNorm3d:?nnq.InstanceNorm3d,
????nn.LayerNorm:?nnq.LayerNorm,
????nn.LeakyReLU:?nnq.LeakyReLU,
????nn.Linear:?nnq.Linear,
????nn.ReLU6:?nnq.ReLU6,
????#?Wrapper?Modules:
????nnq.FloatFunctional:?nnq.QFunctional,
????#?Intrinsic?modules:
????nni.BNReLU2d:?nniq.BNReLU2d,
????nni.BNReLU3d:?nniq.BNReLU3d,
????nni.ConvReLU1d:?nniq.ConvReLU1d,
????nni.ConvReLU2d:?nniq.ConvReLU2d,
????nni.ConvReLU3d:?nniq.ConvReLU3d,
????nni.LinearReLU:?nniq.LinearReLU,
????nniqat.ConvBn1d:?nnq.Conv1d,
????nniqat.ConvBn2d:?nnq.Conv2d,
????nniqat.ConvBnReLU1d:?nniq.ConvReLU1d,
????nniqat.ConvBnReLU2d:?nniq.ConvReLU2d,
????nniqat.ConvReLU2d:?nniq.ConvReLU2d,
????nniqat.LinearReLU:?nniq.LinearReLU,
????#?QAT?modules:
????nnqat.Linear:?nnq.Linear,
????nnqat.Conv2d:?nnq.Conv2d,
}?
在完成模型量化后，我們就要考慮量化模型的推理了。其實(shí)量化模型的推理與浮點(diǎn)模型的推理沒什么本質(zhì)區(qū)別，最大的區(qū)別有這么兩點(diǎn)：
量化節(jié)點(diǎn)插入：需要在網(wǎng)絡(luò)的forward里面插入QuantStub與DeQuantSub兩個(gè)節(jié)點(diǎn)。一個(gè)非常簡(jiǎn)單的參考示例，摘自torchvision.model.quantization.resnet.py。
class?QuantizableResNet(ResNet):
????def?__init__(self,?*args,?**kwargs):
????????super(QuantizableResNet,?self).__init__(*args,?**kwargs)
????????self.quant?=?torch.quantization.QuantStub()
????????self.dequant?=?torch.quantization.DeQuantStub()
????def?forward(self,?x):
????????x?=?self.quant(x)
????????#?Ensure?scriptability
????????#?super(QuantizableResNet,self).forward(x)
????????#?is?not?scriptable
????????x?=?self._forward_impl(x)
????????x?=?self.dequant(x)
????????return?x
op替換：需要將模型中的Add、Concat等操作替換為支持量化的FloatFunctional，可參考如下示例。
class?QuantizableBasicBlock(BasicBlock):
????def?__init__(self,?*args,?**kwargs):
????????super(QuantizableBasicBlock,?self).__init__(*args,?**kwargs)
????????self.add_relu?=?torch.nn.quantized.FloatFunctional()
????def?forward(self,?x):
????????identity?=?x
????????out?=?self.conv1(x)
????????out?=?self.bn1(out)
????????out?=?self.relu(out)
????????out?=?self.conv2(out)
????????out?=?self.bn2(out)
????????if?self.downsample?is?not?None:
????????????identity?=?self.downsample(x)
????????out?=?self.add_relu.add_relu(out,?identity)
????????return?out
參考文章
如何使用PyTorch的量化功能？ PyTorch模型量化工具學(xué)習(xí) Pytorch實(shí)現(xiàn)卷積神經(jīng)網(wǎng)絡(luò)訓(xùn)練量化
個(gè)人微信（如果沒有備注不拉群！） 請(qǐng)注明：地區(qū)+學(xué)校/企業(yè)+研究方向+昵稱
下載1：何愷明頂會(huì)分享
在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù)：何愷明，即可下載。總共有6份PDF，涉及 ResNet、Mask RCNN等經(jīng)典工作的總結(jié)分析
下載2：終身受益的編程指南：Google編程風(fēng)格指南
在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù)：c++，即可下載。歷經(jīng)十年考驗(yàn)，最權(quán)威的編程規(guī)范！
下載3 CVPR2020 在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù)：CVPR2020，即可下載1467篇CVPR?2020論文
覺得不錯(cuò)就點(diǎn)亮在看吧
