Pytorch量化入門之超分量化(一)

 點擊上方AI算法與圖像處理”,選擇加"星標"或“置頂

重磅乾貨,第一時間送達

來源:AIWakler

最近Happy在嘗試進行圖像超分的INT8量化,發現:pytorch量化裏面的坑真多,遠不如TensorFlow的量化好用。不過花了點時間終於還是用pytorch把圖像超分模型完成了量化,以EDSR爲例,模型大小73%,推理速度提升40%左右(PC端),視覺效果幾乎無損,定量指標待補充。有感於網絡上介紹量化的博客一堆,但真正有幫助的較少,所以Happy會盡量以圖像超分爲例提供一個完整的可復現的量化示例(分兩章內容進行)。

背景

量化在不同領域有不同的定義,而在深度學習領域,量化有兩個層面的意義:(1) 存儲量化,即更少的bit來存儲原本需要用浮點數(一般爲FP32)存儲的tensor;(2) 計算量化,即用更少的bit來完成原本需要基於浮點數(一般爲FP32,FP16現在也是常用的一種)完成的計算。量化一般有這樣兩點好處:

  • 更小的模型體積,理論上減少爲FP32模型的75%左右,從筆者不多的經驗來看,往往可以減少73%;
  • 更少的內存訪問與更快的INT8計算,從筆者的幾個簡單嘗試來看,一般可以加速40%左右,這個還會跟平臺相關。

對於量化後模型而言,其部分或者全部tensor(與量化方式、量化op的支持程度有關)將採用INT類型進行計算,而非量化前的浮點類型。量化對於底層的硬件支持、推理框架等要求還是比較高的,目前X86CPU,ARMCPU,Qualcomm DSP等主流硬件對量化都提供了支持;而NCNN、MACE、MNN、TFLite、Caffe2、TensorRT等推理框架也都對量化提供了支持,不過不同框架的支持度還是不太一樣,這個就不細說了,感興趣的同學可以自行百度一下。

筆者主要用Pytorch進行研發,所以花了點精力對其進行了一些研究&嘗試。目前Pytorch已經更新到了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:動態量化,推理過程中的量化,這種量化方式常見諸於NLP領域,在CV領域較少應用;
  • Post Training Static Quantization:靜態量化,訓練後靜態量化,這是CV領域應用非常多的一種量化方式;
  • Quantization Aware Training:感知量化,邊訓練邊量化,一種比靜態量化更優的量化方式,但量化時間會更長,但精度幾乎無損。

注:筆者主要關注CV領域,所以本文也將主要介紹靜態量化與感知量化這種方式。

Tensor量化

要實現量化,那麼就不可避免會涉及到tensor的量化,一般來說,量化公式可以描述如下:

目前Pytorch中的tensor支持int8/uint8/int32等類型的數據,並同時scale、zero_point、quantization_scheme等量化信息。這裏,我們給出一個tensor量化的簡單示例:

x = torch.rand(33)
print(x)
x = torch.quantize_per_tensor(x, scale=0.2, zero_point=3, dtype=torch.quint8)
print(x)
print(x.int_repr())

一個參考輸出如下所示:

image-20210202175857986

注1:藍框爲原始的浮點數據,紅框爲tensor的量化信息,綠框則對應了量化後的INT8數值。

注2:量化不可避免會出現精度損失,這個損失與scale、zero_point有關。

在量化方面,Tensor一般有兩種量化模式:per tensor與per channel。對於PerTensor而言,它的所有數值都按照相同方式進行scale和zero_point處理;而對於PerChannel而言,它有多種不同的scale和zero_point參數,這種方式的量化精度損失更少。

Post Training Static Quantization

靜態量化一般有兩種形式:(1) 僅weight量化;(2) weight與activation同時量化。對於第一種“僅weight量化”而言,只針對weight量化可以使得模型參數所佔內存顯著減小,但在實際推理過程中仍需要轉換成浮點數進行計算;而第二種“weight與activation同時量化”則不僅對weight進行量化,還需要結合校驗數據進行activation的量化。第一種的量化非常簡單,這裏略過,本文僅針對第二種方式進行介紹。

Pytorch的靜態量化一把包含五個步驟:

  • fuse_model:該步驟用來對可以融合的op進行融合,比如Conv與BN的融合、Conv與ReLU的融合、Conv與BN以及ReLU的融合、Linear與BN的融合、Linear與BN以及ReLU的融合。目前Pytorch已經內置的融合code:
fuse_modules(model, modules_to_fuse, inplace=False, fuser_func=fuse_known_modules, fuse_custom_config_dict=None)

在完成融合後,第一個op將被替換會融合後的op,而其他op則會替換爲nn.Identity

  • qconfig:該步驟用於設置用於模型量化的方式,它將插入兩個observer,一個用於監測activation,一個用於監測weight。考慮到推理平臺的不同,pytorch提供了兩種量化配置:針對x86平臺的 fbgemm以及針對arm平臺的 qnnpack

不同平臺的量化配置方式存在些微的區別,大概如下:

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:該步驟用於給每個支持量化的模塊插入Observer,用於收集數據並進行量化數據分析。以activation爲例,它將根據所喂入數據統計min_val與max_val,一般觀察幾個次迭代即可,然後根據所觀察到數據進行統計分析得到scale與zero_point。
  • Feed Data:爲了更好的獲得activation的量化參數信息,我們需要一個合適大小的校驗數據,並將其送入到前述模型中。這個就比較簡單了,就按照模型驗證方式往裏面送數據就可以了。
  • Convert:在完成前述四個步驟後,接下來就需要將完成量化的模型轉換爲量化後模型了,這個就比較簡單了,通過如下命令即可。
torch.quantization.convert(model, inplace=True)

該過程本質上就是用量化OP替換模型中的費量化OP,比如用nnq.Conv2d替換nn.Conv2d, nnq.ConvReLU2d替換nni.ConvReLU2d(注:這是Conv與ReLU的合併)。之前的量化op以及對應的被替換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,

在完成模型量化後,我們就要考慮量化模型的推理了。其實量化模型的推理與浮點模型的推理沒什麼本質區別,最大的區別有這麼兩點:

  • 量化節點插入:需要在網絡的forward裏面插入QuantStub與DeQuantSub兩個節點。一個非常簡單的參考示例,摘自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





參考文章

  1. 如何使用PyTorch的量化功能?
  2. PyTorch模型量化工具學習
  3. Pytorch實現卷積神經網絡訓練量化

   
       
       
       
個人微信(如果沒有備註不拉羣!
請註明: 地區+學校/企業+研究方向+暱稱



下載1:何愷明頂會分享


AI算法與圖像處理」公衆號後臺回覆:何愷明,即可下載。總共有6份PDF,涉及 ResNet、Mask RCNN等經典工作的總結分析


下載2:終身受益的編程指南:Google編程風格指南


AI算法與圖像處理」公衆號後臺回覆:c++,即可下載。歷經十年考驗,最權威的編程規範!



 
     
     
     
下載3 CVPR2020

AI算法與圖像處公衆號後臺回覆: CVPR2020 即可下載1467篇CVPR 2020論文


覺得不錯就點亮在看吧


本文分享自微信公衆號 - AI算法與圖像處理(AI_study)。
如有侵權,請聯繫 [email protected] 刪除。
本文參與“OSC源創計劃”,歡迎正在閱讀的你也加入,一起分享。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章