前言

手寫 AI 推出的全新 TensorRT 模型量化實(shí)戰(zhàn)課程，鏈接。記錄下個(gè)人學(xué)習(xí)筆記，僅供自己參考。

該實(shí)戰(zhàn)課程主要基于手寫 AI 的 Latte 老師所出的 TensorRT下的模型量化，在其課程的基礎(chǔ)上，所整理出的一些實(shí)戰(zhàn)應(yīng)用。

本次課程為 YOLOv7 量化實(shí)戰(zhàn)第一課，主要介紹 TensorRT 量化工具箱 pytorch_quantization。

課程大綱可看下面的思維導(dǎo)圖

1. 課程介紹

什么是模型量化呢？那我們都知道模型訓(xùn)練的時(shí)候是使用的 float32 或 float16 的浮點(diǎn)數(shù)進(jìn)行運(yùn)算，這樣模型能保持一個(gè)比較好的效果，但浮點(diǎn)數(shù)在提升計(jì)算精度的同時(shí)也導(dǎo)致了更多的計(jì)算量以及存儲(chǔ)空間的占用。

由于在模型推理階段我們并不需要進(jìn)行梯度反向傳播，因此我們不需要那么高的計(jì)算精度，這時(shí)可以將高精度的模型參數(shù)映射到低精度上，可以降低運(yùn)算量提高推理速度。

將模型從高精度運(yùn)算轉(zhuǎn)換到低精度運(yùn)算的過(guò)程就叫做模型量化

量化的過(guò)程與數(shù)據(jù)的分布有關(guān)，當(dāng)數(shù)據(jù)分布比較均勻的時(shí)候，高精度 float 向低精度 int 進(jìn)行映射時(shí)就會(huì)將空間利用得比較充分，如果數(shù)據(jù)分布不均勻就會(huì)浪費(fèi)很大的表示空間。

量化又分為飽和量化和非飽和量化，如果直接將量化閾值設(shè)置為 $|x_{\text{max}}|$，此時(shí) INT8 的表示空間沒(méi)有被充分的利用，這是非飽和量化

如果選擇了一個(gè)比較合適的閾值，舍棄那些超出范圍的數(shù)值，再進(jìn)行量化，那這種量化因?yàn)槌浞掷?INT8 的表示空間因此也被稱為飽和量化。

模型量化及其意義可以總結(jié)為：

• 模型量化是指將神經(jīng)網(wǎng)絡(luò)的浮點(diǎn)轉(zhuǎn)換為定點(diǎn)
• 模型量化主要意義就是加快模型端側(cè)的推理速度，并降低設(shè)備功耗和減少存儲(chǔ)空間，工業(yè)界一般只使用 INT8 量化模型。

本系列實(shí)戰(zhàn)課程需要大家具備一些基本的量化知識(shí)，如果對(duì)模型量化知識(shí)模糊的看官的可以先觀看 TensorRT下的模型量化 課程。

2. pytorch_quantization

我們先對(duì) TensorRT 的量化工具箱 pytorch_quantization 做一個(gè)簡(jiǎn)單的介紹

它的安裝指令如下：

pip install pytorch-quantization --extra-index-url https://pypi.ngc.nvidia.com

要求：torch >= 1.9.1，Python >= 3.7， GCC >= 5.4

在博主之前學(xué)習(xí)的過(guò)程中，發(fā)現(xiàn) pytorch 的版本和 pytorch_quantization 的版本如果不適配可能會(huì)導(dǎo)致一些問(wèn)題。

目前博主的軟件版本是：pytorch==2.0.1，pytorch_quantization==2.1.3

我們下面介紹下 pytorch_quantization 工具庫(kù)中的一些函數(shù)、類和模塊

2.1 initialize函數(shù)

首先是 quant_modules 模塊中的 initialize() 函數(shù)，它的使用如下：

import torchvision
from pytorch_quantization import quant_modules

quant_modules.initialize()  # quant_modules 初始化,自動(dòng)為模型插入量化節(jié)點(diǎn)
model = torchvision.models.resnet50()   # 加載 resnet50 模型
# model 是帶有量化節(jié)點(diǎn)的模型

它的作用是初始化量化相關(guān)的設(shè)置和一些參數(shù)，因此我們需要在量化之前調(diào)用它。因?yàn)椴煌愋偷纳窠?jīng)網(wǎng)絡(luò)層如 Conv、Linear、Pool 等等，它們所需要的量化方法是不同的，例如某個(gè)網(wǎng)絡(luò)層當(dāng)中的校準(zhǔn)方法可能用的是 Max，也有可能用的是直方圖，那這都可以在我們量化之前通過(guò) initialize 來(lái)進(jìn)行一個(gè)設(shè)置。

initialize 還有一個(gè)作用，那就是將模型中的 torch 網(wǎng)絡(luò)層替換為相應(yīng)的 quant 量化層，如下所示：

torch.nn.Conv2d     ->  quant_modules.quant_nn.Conv2d
torch.nn.Linear     ->  quant_modules.quant_nn.Linear
torch.nn.MaxPool2d  ->  quant_modules.quant_nn.MaxPool2d

也就是會(huì)把 torch 中對(duì)應(yīng)的算子轉(zhuǎn)換為相應(yīng)的量化版本。

總的來(lái)說(shuō)，initialize 用于在量化模型之前，對(duì)量化過(guò)程進(jìn)行必要的配置和準(zhǔn)備工作以確保量化操作時(shí)按照我們所需要的方式進(jìn)行，這樣的話有助于提高量化模型的性能。

在我們調(diào)用 initialize 之后，我們的模型結(jié)構(gòu)會(huì)插入 FQ 節(jié)點(diǎn)，也就是 fake 算子，如下圖所示：

那在之后的代碼講解部分我們會(huì)清晰的觀察到在調(diào)用 initialize 前后模型結(jié)構(gòu)的一些變化。

2.2 tensor_quant模塊

然后是 tensor_quant 模塊，它的使用如下：

from pytorch_quantization import tensor_quant

tensor_quant.fake_tensor_quant()
tensor_quant.tensor_quant()

tensor_quant 模塊負(fù)責(zé)進(jìn)行張量數(shù)據(jù)的量化操作。那在模型量化過(guò)程中我們有兩種量化方式：

• 模型 weights 的量化：對(duì)于權(quán)重的量化我們是對(duì)權(quán)重的每個(gè)通道進(jìn)行量化，比如一個(gè) Conv 層的通道數(shù)是 32，這意味著 32 個(gè)通道數(shù)的每個(gè)通道都有一個(gè)對(duì)應(yīng)的 scale 值去進(jìn)行量化。
• 模型 inputs/activate 的量化：而對(duì)于輸入或者激活函數(shù)數(shù)值而言，它們的量化是對(duì)每個(gè)張量進(jìn)行量化，也就是說(shuō)整個(gè) Tensor 數(shù)據(jù)都是用同一個(gè) scale 值進(jìn)行量化

具體見下圖：

在上面的圖中我們可以清楚的看到右邊是我們的輸入量化，inputs 的量化 scale 只有一個(gè)，而左邊是我們的權(quán)重量化，weights 的量化 scale 有 32 個(gè)，這是因?yàn)?Conv 的通道數(shù)是 32，它有 32 個(gè) channel，每個(gè) channel 對(duì)應(yīng)一個(gè) scale。

下面的代碼使用了 tensor_quant 模塊中的函數(shù)對(duì)張量進(jìn)行量化：

fake_quant_x   = tensor_quant.fake_tensor_quant(x, x.abs().max) # Q 和 DQ 節(jié)點(diǎn)組成了 Fake 算子
quant_x, scale = tensor_quant.tensor_quant(x, x.abs().max())    # Q 節(jié)點(diǎn)的輸出和 scale 值

我們先來(lái)看看 tensor_quant 中的兩個(gè)函數(shù)

• tensor_quant.fake_tensor_quant
  • 這個(gè)函數(shù)通常用于模擬量化的過(guò)程，而不是實(shí)際上執(zhí)行量化，也就是我們通常說(shuō)的偽量化
  • 偽量化（Fake Quantization）是一種在訓(xùn)練過(guò)程中模擬量化效果的技術(shù)，但在內(nèi)部仍然保持使用浮點(diǎn)數(shù)。
  • 這樣做的目的是使模型適應(yīng)量化帶來(lái)的精度損失，從而在實(shí)際進(jìn)行量化時(shí)能夠保持較好的性能。
• tensor_quant.tensor_quant
  • 這個(gè)函數(shù)用于實(shí)際對(duì)張量進(jìn)行量化，它將輸入的浮點(diǎn)數(shù)張量轉(zhuǎn)換為定點(diǎn)數(shù)的表示（比如從 floa32 轉(zhuǎn)換為 int8）
  • 這個(gè)過(guò)程涉及確定量化的比例因子 scale 和零點(diǎn) zero-point，然后應(yīng)用這些參數(shù)將浮點(diǎn)數(shù)映射到量化的整數(shù)范圍內(nèi)。

在上面的代碼中，x 是我們的輸入數(shù)據(jù)，x.abs().Max 代表我們使用基于 Max 的對(duì)稱量化方法進(jìn)行量化，函數(shù)的輸出 fake_quant_x 是經(jīng)過(guò)偽量化處理的張量，它看起來(lái)像是被量化了，但實(shí)際上仍然是浮點(diǎn)數(shù)。

tensor_quant 函數(shù)的輸出 quant_x 是我們經(jīng)過(guò)實(shí)際 int 量化處理后得到的 int 類型的張量，scale 則是我們用于量化過(guò)程中的比例因子。

2.3 TensorQuantizer類

下面我們來(lái)看看將量化后的模型導(dǎo)出要做哪些操作，實(shí)際上我們需要使用到 nn 模塊中的 TensorQuantizer，它的使用如下：

from pytorch_quantization import nn as quant_nn

quant_nn.TensorQuantizer.use_fb_fake_quant = True   # 模型導(dǎo)出時(shí)將一個(gè) QDQ 算子導(dǎo)出兩個(gè) op

其中 pytorch_quantizaiton 的 nn 模塊提供了量化相關(guān)的神經(jīng)網(wǎng)絡(luò)層和工具，大家可以類比于 pytorch 中的 nn 模塊。而 TensorQuantizer 是一個(gè)用于張量量化的工具類，use_fb_fake_quant 是它的一個(gè)類屬性，用于控制量化過(guò)程中偽量化的行為。

我們將 use_fb_fake_quant 設(shè)置為 True 表明我們?cè)趯?dǎo)出量化模型時(shí)，希望將量化和反量化操作過(guò)程作為兩個(gè)單獨(dú)的 op 算子來(lái)導(dǎo)出，如下圖所示：

可以看到上圖中的紅色框部分，導(dǎo)出的量化模型中包含 QuantizeLinear 和 DequantizeLinear 兩個(gè)模塊，對(duì)應(yīng)我們的量化和反量化兩個(gè) op。

在我們將 use_fb_fake_quant 設(shè)置為 True 的時(shí)候，它會(huì)調(diào)用的是 pytorch 模塊中的兩個(gè)函數(shù)，如下：

torch.fake_quantize_per_tensor_affine
torch.fake_quantize_per_channel_affine

這兩個(gè)函數(shù)會(huì)導(dǎo)出我們之前量化的操作，值得注意的是，在模型導(dǎo)出和模型前向階段的量化操作并不是使用 tensor_quant 模塊中的函數(shù)來(lái)實(shí)現(xiàn)的，而是使用 torch 中上述兩個(gè)函數(shù)來(lái)實(shí)現(xiàn)，這樣做是因?yàn)楦菀邹D(zhuǎn)化成相應(yīng) 的 tensorRT 的一個(gè)操作符，以便我們后續(xù)的部署。在模型訓(xùn)練階段，我們則是調(diào)用 tensor_quant 函數(shù)插入 fake 算子來(lái)進(jìn)行量化的，大家需要了解到在模型訓(xùn)練和前向階段調(diào)用的函數(shù)的不同。

在 Torch-TesorRT 內(nèi)部，fake_quantize_per_*_affine 會(huì)被轉(zhuǎn)換為 QuantizeLayer 和 DequantizerLayer，也就是我們上面導(dǎo)出 ONNX 模型的兩個(gè) op 算子。

從上圖中我們能清晰的看出在模型訓(xùn)練的時(shí)候和模型導(dǎo)出的時(shí)候 Q/DQ 節(jié)點(diǎn)所發(fā)生的一個(gè)變化，在模型訓(xùn)練的時(shí)候，我們是通過(guò) tensor_quant 來(lái)插入 fake 算子來(lái)實(shí)現(xiàn)量化的，而在模型訓(xùn)練完成后導(dǎo)出 ONNX 時(shí)，我們是需要將 use_fb_fake_quant 置為 True，它會(huì)調(diào)用 torch 中的函數(shù)將 fake 算子的節(jié)點(diǎn)導(dǎo)出成 Q 和 DQ 兩個(gè)模塊。

2.4 QuantDescriptor類

接下來(lái)我們?cè)賮?lái)看下 QuantDescriptor 類，它的使用如下：

import torch
import pytorch_quantization.nn as quant_nn
from pytorch_quantization.tensor_quant import QuantDescriptor

# 自定義層的量化
class QuantMultiAdd(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self._input_quantizer  = quant_nn.TensorQuantizer(QuantDescriptor(
                                 num_bits=8, calib_method="histogram"))
        self._weight_quantizer = quant_nn.TensorQuantizer(QuantDescriptor(
                                 num_bits=8, axis=(1), calib_method="histogram"))
    def forward(self, w, x, y):
        return self._weight_quantizer(w) * self._input_quantizer(x) + self._input_quantizer(y)

QuantDescriptor 類主要是用于配置量化的描述符，包括量化的位數(shù)，量化的方法等等。在上面的代碼中，我們創(chuàng)建了一個(gè)自定義的量化層，該層對(duì)權(quán)重和輸入進(jìn)行量化，并執(zhí)行加權(quán)乘法和加法操作

• 我們先創(chuàng)建了兩個(gè) TensorQuantizer 實(shí)例，一個(gè)是 _input_quantizer 用于輸入量化，另一個(gè)是 _weight_quantizer 用于權(quán)重量化
• 我們使用 QuantDescriptor 來(lái)描述量化的參數(shù)，對(duì)于這兩個(gè)量化器，都使用了 8bit 量化，量化的校準(zhǔn)方法都設(shè)置為直方圖校準(zhǔn)

也就是說(shuō)，我們使用 QuantDescriptor 可以實(shí)現(xiàn)自定義層的量化操作，在后續(xù)代碼介紹的時(shí)候會(huì)使用到這個(gè)類。

2.5 calib模塊

我們?cè)賮?lái)看下 pytorch_quantization 中的校準(zhǔn)模塊 calib，它的使用如下：

from pytorch_quantization import calib

if isinstance(module._calibrator, calib.MaxCalibrator):
    module.load_calib_amax()

calib 校準(zhǔn)模塊包含 MaxCalibrator 和 HistogramCalibrator 兩個(gè)校準(zhǔn)類，其中 MaxCalibrator 用于執(zhí)行最大值校準(zhǔn)，在我們的量化訓(xùn)練中，我們通常會(huì)確定每個(gè)張量的一個(gè)動(dòng)態(tài)范圍，也就是它們的最大值和最小值，Max 方法通過(guò)跟蹤張量的最大值來(lái)執(zhí)行標(biāo)定工作，以便在量化推理時(shí)能將其映射到 int 整數(shù)范圍之內(nèi)。

而對(duì)于 Histogram 直方圖校準(zhǔn)方法則是通過(guò)收集和分析張量值的直方圖來(lái)確定我們的動(dòng)態(tài)范圍，這種方法可以更準(zhǔn)確地估計(jì)張量值的一個(gè)分布，并且更好地適應(yīng)不同數(shù)據(jù)分布的情況。

這兩種校準(zhǔn)方法在模型量化中都有它們各自的優(yōu)勢(shì)，具體選擇哪種校準(zhǔn)方法主要取決于我們具體的應(yīng)用場(chǎng)景和數(shù)據(jù)分布的情況，我們通常是根據(jù)數(shù)據(jù)分布和量化的需求來(lái)選擇合適的校準(zhǔn)方法，以確保量化后的模型在推理時(shí)能保持一個(gè)比較好的準(zhǔn)確性。

以上就是關(guān)于 pytorch_quantization 中的函數(shù)、類和模塊的簡(jiǎn)單介紹。

總結(jié)

本次課程介紹了 pytorch_quantization 量化工具以及其中的一些函數(shù)、類和模塊。在我們量化之前需要調(diào)用 initialize 函數(shù)來(lái)初始化量化相關(guān)的一些設(shè)置和參數(shù)。接著我們會(huì)使用 tensor_quant 模塊來(lái)對(duì)張量數(shù)據(jù)進(jìn)行實(shí)際的量化，而在量化完成后導(dǎo)出時(shí)我們需要將 TensorQuantizer 類中的屬性 usb_fb_fake_quant 設(shè)置為 True，使得導(dǎo)出的量化模型包含 Q、DQ 兩個(gè)模塊。這是因?yàn)樵谀Ｐ陀?xùn)練階段和前向、導(dǎo)出階段的量化操作調(diào)用的函數(shù)是不同的，訓(xùn)練階段是通過(guò) tensor_quant 函數(shù)插入 fake 算子來(lái)量化的，而導(dǎo)出階段是 torch 中的兩個(gè)函數(shù)來(lái)實(shí)現(xiàn)的。

在量化過(guò)程中我們還會(huì)使用 QuantDescriptor 來(lái)配置量化的一些參數(shù)，包括量化位數(shù)、量化方法等等，最后我們簡(jiǎn)單介紹了 Calib 校準(zhǔn)模塊，它包含 Max 和 Histogram 兩種校準(zhǔn)方法。

下節(jié)我們正式進(jìn)入 YOLOv7-PTQ 量化的學(xué)習(xí)??文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-732547.html

到了這里，關(guān)于TensorRT量化實(shí)戰(zhàn)課YOLOv7量化：pytorch_quantization介紹的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！