HCHUNGW的部落格

12.7 PTQ 量化實踐

前言

上一節介紹了模型量化的概念，本節開始進行PTQ的實踐，並對比不同動態範圍計算方法的差別（max, entropy, mse, pertentile）。

由於量化是對已經訓練好的模型進行量化並評估在測試集上的精度，因此量化（PTQ和QAT）需要依託一個完成的模型訓練專案開展，這裡採用第八章第一節的分類任務。

PTQ代碼實踐將從以下幾個部分進行：pytorch-quantization庫介紹、推理統計、校準、量化、評估模型準確率和TensorRT推理效率。

pytorch_quantization 工具庫

安裝：

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

或者（pytorch 1.12版本時推薦用 pytorch-quantization==2.1.2 ）

pip install pytorch-quantization==2.1.2

Copy

pytorch_quantization是由NVIDIA官方提供的PyTorch量化庫，用於將PyTorch模型量化為低比特形式（如Int8）。

pytorch_quantization 庫使用非常方便，其PTQ步驟分四步:

第一步：構建具備量化模組的nn.Module，在模型定義之前替換pytorch官方nn.Module的網路層。如nn.conv2替換為quant_nn.QuantConv2d。

Quantxxx這樣的網路層中會在原功能基礎上，添加一些元件，這些元件將在校準時記錄資料、並存儲scale和Z值。

第二步：正常構建模型，並執行前向推理。

第三步：執行動態範圍計算，並計算scale和Z值。

第四步：匯出ONNX模型

通過以上流程可知道，僅第一步中pytorch_quantization 庫中自訂的量化功能模組需要深入瞭解，其餘都是常規流程。

因此，下面簡單介紹pytorch_quantization中的核心元件。

量化函數——Quantization function

提供兩個量化功能函數 tensor_quant 和 fake_tensor_quant。

• tensor_quant(inputs, amax, num_bits=8, output_dtype=torch.float, unsigned=False)
•  
  • 功能：對輸入張量進行量化，輸出量化後的張量
• fake_tensor_quant(inputs, amax, num_bits=8, output_dtype=torch.float, unsigned=False)
•  
  • 輸入張量進行偽量化，先量化，後反量化，偽量化輸出的張量與原張量存在量化誤差，可用於理解量化誤差的產生。

from pytorch_quantization import tensor_quant

# Generate random input. With fixed seed 12345, x should be

# tensor([0.9817, 0.8796, 0.9921, 0.4611, 0.0832, 0.1784, 0.3674, 0.5676, 0.3376, 0.2119])

torch.manual_seed(12345)

x = torch.rand(10)

# fake quantize tensor x. fake_quant_x will be

# tensor([0.9843, 0.8828, 0.9921, 0.4609, 0.0859, 0.1797, 0.3672, 0.5703, 0.3359, 0.2109])

fake_quant_x = tensor_quant.fake_tensor_quant(x, x.abs().max())

# quantize tensor x. quant_x will be

# tensor([126., 113., 127.,  59.,  11.,  23.,  47.,  73.,  43.,  27.])

# with scale=128.0057

quant_x, scale = tensor_quant.tensor_quant(x, x.abs().max())

Copy

量化描述符和量化器——Descriptor and quantizer

QuantDescriptor 用於描述該採用何種量化方式，例如 QUANT_DESC_8BIT_PER_TENSOR、 QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL。

這裡的pre tensor ， per channel是一個新知識點，對於activation通常用per tensor， 對於卷積核通常用per channel。細節不展開，感興趣可以閱讀《A White Paper on Neural Network Quantization》的2.4.2 Per-tensor and per-channel quantization

描述符是用於初始化量化器quantizer的，量化器可接收張量，輸出量化後的張量。完成量化後，在量化器內部會記錄相應的scale和Z值。

from pytorch_quantization.tensor_quant import QuantDescriptor

from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer

quant_desc = QuantDescriptor(num_bits=4, fake_quant=False, axis=(0), unsigned=True)

quantizer = TensorQuantizer(quant_desc)

torch.manual_seed(12345)

x = torch.rand(3, 4, 2)

quant_x = quantizer(x)

print(x)

print(quant_x)

print(quantizer.scale)  # 可以得到3個scale，可知道是按照第一個維度切分張量。因為描述符設置了 axis=(0)

Copy

量化模組——Quantized module

前面提到PTQ量化第一步是將pq庫的Quantxxx模組替換掉pytorch的nn.Module，例如nn.conv2d變為quant_nn.Conv2d。

在實際量化中主要有Linear和Conv兩大類（nn.Module可以回顧chapter-4），下面對比pytorch和pq庫的網路層創建區別。

from torch import nn

from pytorch_quantization import tensor_quant

import pytorch_quantization.nn as quant_nn

# pytorch's module

fc1 = nn.Linear(in_features, out_features, bias=True)

conv1 = nn.Conv2d(in_channels, out_channels, kernel_size)

# quantized version

quant_fc1 = quant_nn.Linear(

    in_features, out_features, bias=True,

    quant_desc_input=tensor_quant.QUANT_DESC_8BIT_PER_TENSOR,

    quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW)

quant_conv1 = quant_nn.Conv2d(

    in_channels, out_channels, kernel_size,

    quant_desc_input=tensor_quant.QUANT_DESC_8BIT_PER_TENSOR,

    quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL)

Copy

可以發現，只是多了量化描述符的配置，根據量化物件（上一節提到，主要是weight， bias， activation），通常需要設定不同的量化方式。

例如，對於activation用per tensor， 對於linear的權重用per row，對於conv的權重用pre channel。

ResNet50 PTQ

模型權重傳入百度雲盤，供下載

下面基於chapter-8/01_classification中的項目進行量化，首先訓練一個resnet50，得到的accuracy是94.39。模型權重下載-提取碼：jhkm

# 也可以重新訓練 resnet50

nohup python train_main.py --data-path ../data/chest_xray --batch-size 64 --workers 8 --lr 0.01 --lr-step-size 20 --epochs 50 --model resnet50 > ./log.log 2>&1 &

Copy

接下來對它進行量化，並對比accuracy掉點情況以及推理加速情況。

配套代碼位於這裡

流程分析

PTQ整體流程在前言以及介紹，這裡從代碼角度介紹核心流程

# 第一步，初始化pq，將pytorch的模組替換掉

quant_modules.initialize()    # 替換torch.nn的常用層，變為可量化的層

# 第二步，創建模型、載入訓練權重

model = get_model(args, logger, device)

# 第三步，前向推理，統計activation啟動值分佈

collect_stats(model, train_loader, num_batches=args.num_data)  # 設置量化模組開關，並推理，同時統計啟動值

# 第四步，根據動態範圍方法計算scale、Z值，獲得量化後的模型

compute_amax(model, method=args.ptq_method)                     # 計算上限、下限，並計算scale 、Z值

# 第五步，採用量化後的模型進行accuracy評估

loss_m_valid, acc_m_valid, mat_valid = \

    utils.ModelTrainer.evaluate(valid_loader, model, criterion, device, classes)

# 第六步，保存成ONNX模型和pth模型

Copy

第一步中，可以進入pytorch_quantization\quant_modules.py查看當前支持的量化層，量化層中會插入量化器，量化器中存儲量化用的資料——scale和Z

# Global member of the file that contains the mapping of quantized modules

_DEFAULT_QUANT_MAP = [_quant_entry(torch.nn, "Conv1d", quant_nn.QuantConv1d),

                      _quant_entry(torch.nn, "Conv2d", quant_nn.QuantConv2d),

                      _quant_entry(torch.nn, "Conv3d", quant_nn.QuantConv3d),

                      _quant_entry(torch.nn, "ConvTranspose1d", quant_nn.QuantConvTranspose1d),

                      _quant_entry(torch.nn, "ConvTranspose2d", quant_nn.QuantConvTranspose2d),

                      _quant_entry(torch.nn, "ConvTranspose3d", quant_nn.QuantConvTranspose3d),

                      _quant_entry(torch.nn, "Linear", quant_nn.QuantLinear),

                      _quant_entry(torch.nn, "LSTM", quant_nn.QuantLSTM),

                      _quant_entry(torch.nn, "LSTMCell", quant_nn.QuantLSTMCell),

                      _quant_entry(torch.nn, "AvgPool1d", quant_nn.QuantAvgPool1d),

                      _quant_entry(torch.nn, "AvgPool2d", quant_nn.QuantAvgPool2d),

                      _quant_entry(torch.nn, "AvgPool3d", quant_nn.QuantAvgPool3d),

                      _quant_entry(torch.nn, "AdaptiveAvgPool1d", quant_nn.QuantAdaptiveAvgPool1d),

                      _quant_entry(torch.nn, "AdaptiveAvgPool2d", quant_nn.QuantAdaptiveAvgPool2d),

                      _quant_entry(torch.nn, "AdaptiveAvgPool3d", quant_nn.QuantAdaptiveAvgPool3d),]

Copy

第二步，載入模型。

第三步，進行前向推理，但在實際推理時需要對兩個開關分別進行控制。

• 量化開關：若打開，會進行量化，這個在校準的時候是不需要的。
• 校準開關：若打開，會在模型前向傳播時，統計資料，這個在校準時是需要的。

因此在推理前有這樣的程式碼片段（同理，推理後需要把兩個開關反過來）：

# Enable calibrators

for name, module in model.named_modules():

    if isinstance(module, quant_nn.TensorQuantizer):

        if module._calibrator is not None:

            module.disable_quant()

            module.enable_calib()

        else:

            module.disable()

Copy

第四步中，調用module.load_calib_amax(**kwargs)，實現scale和Z值的計算。scale和Z值會存在量化器中。

例如resnet50中的conv1層是一個QuantConv2d，其中包含了input_quantizer和weight_quantizer兩個量化器，量化器中存儲了amax和step_size(scale值)

第五步中，量化後resnet50模型推理時，量化層會對輸入資料、權重進行量化，然後再進行運算。

這裡debug進入到layer1的第一個卷積層，觀察Quantconv2d的forward（pytorch_quantization\nn\modules\quant_conv.py）

可以看到，進行運算前，對輸入資料和權重進行量化，量化用到的就是上述提到的量化器，量化器中有scale和Z值。

def forward(self, input):

    quant_input, quant_weight = self._quant(input)

    output = F.conv2d(quant_input, quant_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)

    return output

def _quant(self, input):

    quant_input = self._input_quantizer(input)

    quant_weight = self._weight_quantizer(self.weight)

    return (quant_input, quant_weight)

Copy

具體量化過程，可debug追溯得到如下過程：

quant_weight = self._weight_quantizer(self.weight)

1 ---> 進入量化器中的forward。pytorch_quantization\nn\modules\tensor_quantizer.py

1 ---> 量化器中可實現校準功能、截斷功能、量化功能，是根據3個屬性開關來判斷，這裡也發現了校準資料收集的代碼。

def forward(self, inputs):

    if self._disabled:

        return inputs

    outputs = inputs

    if self._if_calib:

        if self._calibrator is None:

            raise RuntimeError("Calibrator was not created.")

        # Shape is only know when it sees the first tensor

        self._calibrator.collect(inputs)

    if self._if_clip:

        if not self._learn_amax:

            raise RuntimeError("Clip without learning amax is not implemented.")

        outputs = self.clip(inputs)

    if self._if_quant:

        outputs = self._quant_forward(inputs)

    return outputs

2 ---> 進入量化器中的_quant_forward(),

2 ---> _quant_forward中可以實現偽量化、量化兩種操作，這裡是需要進行量化的，因此會執行

2 ---> outputs, self._scale = tensor_quant(inputs, amax, self._num_bits, self._unsigned)

2 ---> tensor_quant就是開篇介紹pytorch_quantization庫的第一個功能，對張量進行量化的函數。

def _quant_forward(self, inputs):

    """Quantized forward pass."""

    if self._learn_amax:

        inputs = self.clip(inputs)

        amax = torch.max(-self.clip.clip_value_min, self.clip.clip_value_max).detach()

    else:

        amax = self._get_amax(inputs)

    if self._fake_quant:

        if not TensorQuantizer.use_fb_fake_quant:

            outputs = fake_tensor_quant(inputs, amax, self._num_bits, self._unsigned, self._narrow_range)

        else:

            if inputs.dtype == torch.half or amax.dtype == torch.half:

                raise Exception("Exporting to ONNX in fp16 is not supported. Please export in fp32, i.e. disable AMP.")

            outputs = self._fb_fake_quant(inputs, amax)

    else:

        outputs, self._scale = tensor_quant(inputs, amax, self._num_bits, self._unsigned)

    return outputs

Copy

完成量化模型評估後，可進行模型保存為ONNX格式，再用netron查看。

可以看到量化的層會插入QDQ節點，這些節點在TensorRT中將會被使用

PTQ 量化實現

執行命令後，可在chapter-8\01_classification\Result\2023-09-26_01-47-40 資料夾下獲得對應的onnx輸出，並可觀察不同方法得到的模型精度

python resnet_ptq.py --mode quantize --num-data 512

Copy

動態範圍計算方法對比

對於resnet_ptq.py，這裡簡單介紹使用方法

• 根據args.ptq_method的有無決定單個動態範圍方法量化， 還是四種量化方法均量化。
• 根據args.mode == 'quantize' 還是 evaluate，決定代碼進行量化，還是進行單純的預訓練模型accuracy評估。（用於對比）
• 其他參數參考get_args_parser()

# 進行四種方法對比。跑512個batch

python resnet_ptq.py --mode quantize --num-data 512

# 單個方法量化

python resnet_ptq.py --mode quantize --ptq-method max --num-data 512

python resnet_ptq.py --mode quantize --ptq-method entropy --num-data 512

python resnet_ptq.py --mode quantize --ptq-method mse --num-data 512

python resnet_ptq.py --mode quantize --ptq-method percentile --num-data 512

Copy

通過五次實驗，發現PTQ後的Acc存在較大方差，具體資料分佈如下箱線圖所示：

{'entropy': [85.74, 89.26, 91.67, 92.79, 93.27],

'max': [83.17, 90.54, 92.15, 92.79, 94.07],

'mse': [89.42, 91.67, 92.63, 92.79, 93.59],

'percentile': [85.58, 87.34, 91.67, 92.63, 92.79]}

Copy

效果最好的是max，除了一次83%的acc之外，其餘效果比較好，並且最高達到94.07%。

建議：方法都試試，實在不穩定，考慮QAT。

PTQ 效率對比

為了觀察PTQ帶來的效率提升，下面進行未量化、int8量化，分別在bs=1，bs=32時的對比。

首先獲得fp32的onnx：

 python resnet_ptq.py --mode onnxexport

Copy

此時在Result\2023-09-26_01-47-40下應當擁有所有onnx模型，可以採用trtexec進行效率測試

trtexec --onnx=resnet_50_fp32_bs1.onnx

trtexec --onnx=resnet_50_fp32_bs32.onnx

trtexec --onnx=resnet_50_ptq_bs1_data-num512_percentile_91.03%.onnx --int8

trtexec --onnx=resnet_50_ptq_bs32_data-num512_percentile_91.03%.onnx --int8 --saveEngine=resnet_ptq_int8.engine

Copy

通過實驗，可知道，時延可下降40%， 輸送量提高70%左右。

小結

本節首先介紹pytorch-quantization庫中三個核心概念，然後梳理PTQ量化步驟，接著實現了resnet50的PTQ量化實驗，最後對量化後的模型在trtexec上進行效率對比。

通過本節實踐可知：

• int8 可有效提高推理速度和輸送量，具體提高比例需結合具體模型、GPU型號而定。本案例時延降低40%左右，輸送量提高70%左右
• PTQ量化掉點不穩定，需多次試驗，或進行逐層分析，或逐層設定動態範圍方法
• 採用pytorch-quantization量化，大體分6步：初始化quanti_modules，載入模型，推理統計，量化獲得s和z，精度評估，模型匯出

下一小節，介紹QAT的實現。