Pytorch模型量化
在深度学习中,量化指的是使用更少的bit来存储原本以浮点数存储的tensor,以及使用更少的bit来完成原本以浮点数完成的计算。这么做的好处主要有如下几点:
- 更少的模型体积,接近4倍的减少;
- 可以更快的计算,由于更少的内存访问和更快的int8计算,可以快2~4倍。
一个量化后的模型,其部分或者全部的tensor操作会使用int类型来计算,而不是使用量化之前的float类型。当然,量化还需要底层硬件支持,x86 CPU(支持AVX2)、ARM CPU、Google TPU、Nvidia Volta/Turing/Ampere、Qualcomm DSP这些主流硬件都对量化提供了支持。
PyTorch对量化的支持目前有如下三种方式:
- Post Training Dynamic Quantization:模型训练完毕后的动态量化
- Post Training Static Quantization:模型训练完毕后的静态量化
- QAT (Quantization Aware Training):模型训练中开启量化
在开始这三部分之前,先介绍下最基础的Tensor的量化
量化基础
量化函数
量化:$$公式1:xq=round(\frac{x}{scale}+zero\_point)$$
反量化:$$公式2:x = (xq-zero\_point)*scale$$
式中,scale是缩放因子,zero_point是零基准,也就是fp32中的零在量化tensor中的值
scale是输入范围与输出范围的比例:$scale=\frac{_max - _min}{q\_max - q\_min}$
其中[min, max]是输入的裁剪范围,即允许输入的边界。[q_min, q_max]是量化输出空间中的范围。对于int8量化,输出范围$q\_max - q\_min \le (2^8-1)$
zero_point充当偏差 以确保输入中的0完美映射到量化空间中的0:$zero\_point=q\_min-\frac{min}{scale}$
校准
选择输入限幅范围的过程称为校准。最简单的技术(也是 PyTorch 中的默认技术)是记录运行过程中的最小值和最大值。TensorRT还使用熵最小化(KL 散度)、均方误差最小化或输入范围的百分位数。
在 PyTorch 中,Observer模块收集输入值的统计信息并计算scale和zero_point。不同的校准方案会产生不同的量化输出,最好凭经验验证哪种方案最适合您的应用程序和架构(稍后会详细介绍)。
import torch from torch.quantization.observer import MinMaxObserver, MovingAverageMinMaxObserver, HistogramObserver C, L = 3, 4 normal = torch.distributions.normal.Normal(0, 1) inputs = [normal.sample((C, L)), normal.sample((C, L))] print(inputs) # [tensor([[-0.0590, 1.1674, 0.7119, -1.1270], # [-1.3974, 0.5077, -0.5601, 0.0683], # [-0.0929, 0.9473, 0.7159, -0.4574]]]), # tensor([[-0.0236, -0.7599, 1.0290, 0.8914], # [-1.1727, -1.2556, -0.2271, 0.9568], # [-0.2500, 1.4579, 1.4707, 0.4043]])] observers = [MinMaxObserver(), # 最小值、最大值 观察者 MovingAverageMinMaxObserver(), # 移动平均最小值、最大值 观察者 HistogramObserver()] for obs in observers: for x in inputs: obs(x) print(obs.__class__.__name__, obs.calculate_qparams()) # MinMaxObserver (tensor([0.0112]), tensor([124], dtype=torch.int32)) # MovingAverageMinMaxObserver (tensor([0.0101]), tensor([139], dtype=torch.int32)) # HistogramObserver (tensor([0.0100]), tensor([106], dtype=torch.int32))
仿射量化和对称量化
仿射或非对称量化方案将输入范围分配给最小和最大观察值。仿射方案通常提供更严格的裁剪范围,并且对于量化非负激活很有用(如果输入张量从不为负,则不需要输入范围包含负值)。范围计算为$\alpha=min(r)$,$\beta=max(r)$。当用于权重张量 [ 3 ] 时,仿射量化会导致计算量更大的推理。
对称量化方案将输入范围集中在 0 附近,无需计算零点偏移量。范围计算为$-\alpha=\beta=max(|max(r)|,|min(r)|)$。对于偏斜信号(如非负激活),这可能会导致量化分辨率不佳,因为裁剪范围包括从未出现在输入中的值(请参见下面的 pyplot)。
图 仿射和对称方案的裁剪范围(紫色)
在 PyTorch 中,我们可以在初始化 Observer 时指定仿射或对称方案。请注意,并非所有观察者都支持这两种方案。
for qscheme in [torch.per_tensor_affine, torch.per_tensor_symmetric]: obs = MovingAverageMinMaxObserver(qscheme=qscheme) for x in inputs: obs(x) print(f"Qscheme: {qscheme} | {obs.calculate_qparams()}") # Qscheme: torch.per_tensor_affine | (tensor([0.0101]), tensor([139], dtype=torch.int32)) # Qscheme: torch.per_tensor_symmetric | (tensor([0.0109]), tensor([128]))
每张量和每通道量化方案
- per-Tensor量化:将层的整个权重张量作为一个整体计算量化参数,相同的裁剪范围应用于层中的所有通道,为每个张量对整个张量使用相同的 qparams(scale和offse)。
- per-Channel量化:将每个通道单独计算量化参数,为每个通道使用一组 qparams(scale和offse)。
对于权重量化,per-Channel 对称量化提供更好的精度;per-tensor 量化表现不佳,可能是由于 BatchNorm 折叠 [3] 跨通道的 Conv 权重差异很大。
from torch.quantization.observer import MovingAveragePerChannelMinMaxObserver obs = MovingAveragePerChannelMinMaxObserver(ch_axis=0) # 分别计算所有' C '通道的qparams for x in inputs: obs(x) print(obs.calculate_qparams()) # (tensor([0.0090, 0.0075, 0.0055]), tensor([125, 187, 82], dtype=torch.int32))
Backend Engine
目前,量化运算符通过fbgemm后端在 x86 机器上运行,或者在 ARM 机器上使用qnnpack。对服务器 GPU 的后端支持(通过 TensorRT 和 cuDNN)即将推出。了解有关将量化扩展到自定义后端的更多信息:RFC-0019。
backend = 'fbgemm' if x86 else 'qnnpack' qconfig = torch.quantization.get_default_qconfig(backend) torch.backends.quantized.engine = backend
QConfig
QConfig NamedTuple 存储观察者和用于量化激活和权重的量化方案。请务必传递 Observer 类(而非实例)或返回 Observer 实例的可调用对象。用于with_args()覆盖默认参数。
my_qconfig = torch.quantization.QConfig( activation=MovingAverageMinMaxObserver.with_args(qscheme=torch.per_tensor_affine), weight=MovingAveragePerChannelMinMaxObserver.with_args(qscheme=torch.qint8) ) # >>>>> # QConfig(activation=functools.partial(<class 'torch.ao.quantization.observer.MovingAverageMinMaxObserver'>, qscheme=torch.per_tensor_affine){}, weight=functools.partial(<class 'torch.ao.quantization.observer.MovingAveragePerChannelMinMaxObserver'>, qscheme=torch.qint8){})
Tensor的量化
为了实现量化,PyTorch 引入了能够表示量化数据的Quantized Tensor,可以存储 int8/uint8/int32类型的数据,并携带有scale、zero_point这些参数。把一个标准的float Tensor转换为量化Tensor的步骤如下:
import torch x = torch.randn(2, 2, dtype=torch.float32) # tensor([[ 0.9872, -1.6833], # [-0.9345, 0.6531]]) # 公式1(量化):xq = round(x / scale + zero_point) # 使用给定的scale和 zero_point 来把一个float tensor转化为 quantized tensor xq = torch.quantize_per_tensor(x, scale=0.5, zero_point=8, dtype=torch.quint8) # tensor([[ 1.0000, -1.5000], # [-1.0000, 0.5000]], size=(2, 2), dtype=torch.quint8, # quantization_scheme=torch.per_tensor_affine, scale=0.5, zero_point=8) print(xq.int_repr()) # 给定一个量化的张量,返回一个以 uint8_t 作为数据类型的张量 # tensor([[10, 5], # [ 6, 9]], dtype=torch.uint8) # 公式2(反量化):xdq = (xq - zero_point) * scale # 使用给定的scale和 zero_point 来把一个 quantized tensor 转化为 float tensor xdq = xq.dequantize() # tensor([[ 1.0000, -1.5000], # [-1.0000, 0.5000]])
xdq和x的值已经出现了偏差的事实告诉了我们两个道理:
- 量化会有精度损失
- 我们随便选取的scale和zp太烂,选择合适的scale和zp可以有效降低精度损失。不信你把scale和zp分别换成scale = 0.0036, zero_point = 0试试
pytorch中的量化
PyTorch 允许您使用几种不同的方法来量化您的模型,具体取决于
- 如果您更喜欢灵活和手动或受限的自动过程(Eager Mode v/s FX Graph Mode)
- 如果用于量化激活(层输出)的 qparams 是为所有输入预先计算的,或者用每个输入重新计算(静态v/s动态),
- 如果 qparams 是在有或没有重新训练的情况下计算的(量化感知训练v/s训练后量化)
在我们正式了解pytorch模型量化前我们再来检查一下pytorch的官方量化是否能满足我们的需求,如果不能,后面的都不需要看了
静态量化 | 动态量化 | |
nn.linear | Y | Y |
nn.Conv1d/2d/3d | Y | N (因为pytorch认为卷积参数来了个太小了,对卷积核进行量化会造成更多损失,所以pytorch选择不量化) |
nn.LSTM | N(LSTM的好像又可以了,官方给出了一个例子,传送门) | Y |
nn.GRU | N | Y |
nn.RNNCell | N | Y |
nn.GRUCell | N | Y |
nn.LSTMCell | N | Y |
nn.EmbeddingBag | Y(激活在fp32) | Y |
nn.Embedding | Y | N |
nn.MultiheadAttention | N | N |
Activations | 大部分支持 | 不变,计算停留在fp32中 |
第二点:pytorch模型的动态量化只量化权重,不量化偏置
Post Training Dynamic Quantization (训练后动态量化)
意思就是对训练后的模型权重执行动态量化,将浮点模型转换为动态量化模型,仅对模型权重进行量化,偏置不会量化。默认情况下,仅对 Linear 和 RNN 变体量化 (因为这些layer的参数量很大,收益更高)。
torch.quantization.quantize_dynamic(model, qconfig_spec=None, dtype=torch.qint8, mapping=None, inplace=False)
参数:
- model:浮点模型
- qconfig_spec:
- 下面的任意一种
- 集合:比如: qconfig_spec={nn.LSTM, nn.Linear} 。罗列 要量化的NN
- 字典: qconfig_spec = {nn.Linear : default_dynamic_qconfig, nn.LSTM : default_dynamic_qconfig}
- 下面的任意一种
- dtype: float16 或 qint8
- mapping:就地执行模型转换,原始模块发生变异
- inplace:将子模块的类型映射到需要替换子模块的相应动态量化版本的类型
返回:动态量化后的模型
我们来吃一个栗子:
# -*- coding:utf-8 -*- # Author:凌逆战 | Never # Date: 2022/10/17 """ 只量化权重,不量化激活 """ import torch from torch import nn class DemoModel(torch.nn.Module): def __init__(self): super(DemoModel, self).__init__() self.conv = nn.Conv2d(in_channels=1,out_channels=1,kernel_size=1) self.relu = nn.ReLU() self.fc = torch.nn.Linear(2, 2) def forward(self, x): x = self.conv(x) x = self.relu(x) x = self.fc(x) return x if __name__ == "__main__": model_fp32 = DemoModel() # 创建一个量化的模型实例 model_int8 = torch.quantization.quantize_dynamic( model=model_fp32, # 原始模型 qconfig_spec={torch.nn.Linear}, # 要动态量化的NN算子 dtype=torch.qint8) # 将权重量化为:float16 \ qint8 print(model_fp32) print(model_int8) # 运行模型 input_fp32 = torch.randn(1,1,2, 2) output_fp32 = model_fp32(input_fp32) print(output_fp32) output_int8 = model_int8(input_fp32) print(output_int8)
输出
DemoModel( (conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1)) (relu): ReLU() (fc): Linear(in_features=2, out_features=2, bias=True) ) DemoModel( (conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1)) (relu): ReLU() (fc): DynamicQuantizedLinear(in_features=2, out_features=2, dtype=torch.qint8, qscheme=torch.per_tensor_affine) ) tensor([[[[-0.5361, 0.0741], [-0.2033, 0.4149]]]], grad_fn=<AddBackward0>) tensor([[[[-0.5371, 0.0713], [-0.2040, 0.4126]]]])
Fx模式
import torch from torch import nn # toy model m = nn.Sequential( nn.Conv2d(2, 64, (8,)), nn.ReLU(), nn.Linear(16,10), nn.LSTM(10, 10)) m.eval() from torch.quantization import quantize_fx qconfig_dict = {"": torch.quantization.default_dynamic_qconfig} # 空键表示应用于所有模块的默认值 model_prepared = quantize_fx.prepare_fx(m, qconfig_dict) model_quantized = quantize_fx.convert_fx(model_prepared)
Post Training Static Quantization (训练后静态量化)
静态量化需要把模型的权重和激活都进行量化,静态量化需要把训练集或者和训练集分布类似的数据喂给模型(注意没有反向传播),然后通过每个op输入的分布 来计算activation的量化参数(scale和zp)——称之为Calibrate(定标),因为静态量化的前向推理过程自始至终都是int计算,activation需要确保一个op的输入符合下一个op的输入。另外你也可以参考 pytorch的训练后静态量化博客
PyTorch会使用以下5步来完成模型的静态量化:
1、fuse_model
合并一些可以合并的layer。这一步的目的是为了提高速度和准确度:
fuse_modules(model, modules_to_fuse, inplace=False, fuser_func=fuse_known_modules, fuse_custom_config_dict=None)
比如给fuse_modules传递下面的参数就会合并网络中的fc、relu:
torch.quantization.fuse_modules(F32Model, [['fc', 'relu']], inplace=True)
一旦合并成功,那么原始网络中的fc就会被替换为新的合并后的module(因为其是list中的第一个元素),而relu(list中剩余的元素)会被替换为nn.Identity(),这个模块是个占位符,直接输出输入。举个例子,对于下面的一个小网络:
import torch from torch import nn class F32Model(nn.Module): def __init__(self): super(F32Model, self).__init__() self.fc = nn.Linear(3, 2,bias=False) self.relu = nn.ReLU(inplace=False) def forward(self, x): x = self.fc(x) x = self.relu(x) return x model_fp32 = F32Model() print(model_fp32) # F32Model( # (fc): Linear(in_features=3, out_features=2, bias=False) # (relu): ReLU() # ) model_fp32_fused = torch.quantization.fuse_modules(model_fp32, [['fc', 'relu']]) print(model_fp32_fused) # F32Model( # (fc): LinearReLU( # (0): Linear(in_features=3, out_features=2, bias=False) # (1): ReLU() # ) # (relu): Identity() # )
modules_to_fuse参数的list可以包含多个item list,或者是submodule的op list也可以,比如:[ ['conv1', 'bn1', 'relu1'], ['submodule.conv', 'submodule.relu']]。有的人会说了,我要fuse的module被Sequential封装起来了,如何传参?参考下面的代码:
torch.quantization.fuse_modules(a_sequential_module, ['0', '1', '2'], inplace=True)
就目前来说,截止目前为止,只有如下的op的顺序才可以 (这个mapping关系就定义在DEFAULT_OP_LIST_TO_FUSER_METHOD中):
- Convolution, BatchNorm
- Convolution, BatchNorm, ReLU
- Convolution, ReLU
- Linear, ReLU
- BatchNorm, ReLU
- ConvTranspose, BatchNorm
2、设置qconfig
qconfig要设置到模型或者Module上。
#如果要部署在x86 server上 model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm') #如果要部署在ARM上 model_fp32.qconfig = torch.quantization.get_default_qconfig('qnnpack')
x86和arm之外目前不支持。
3、prepare
prepare用来给每个子module插入Observer,用来收集和定标数据。
以activation的observer为例,观察输入数据得到 四元组中的 min_val 和 max_val,至少观察个几百个迭代的数据吧,然后由这四元组得到 scale 和 zp 这两个参数的值。
model_fp32_prepared= torch.quantization.prepare(model_fp32_fused)
4、喂数据
这一步不是训练。是为了获取数据的分布特点,来更好的计算activation的 scale 和 zp 。至少要喂上几百个迭代的数据。
#至少观察个几百迭代 for data in data_loader: model_fp32_prepared(data)
5、转换模型
第四步完成后,各个op权重的四元组 (min_val,max_val,qmin, qmax) 中的 min_val, max_val 已经有了,各个op activation的四元组 (min_val,max_val,qmin, qmax) 中的 min_val, max_val 也已经观察出来了。那么在这一步我们将调用convert API:
model_prepared_int8 = torch.quantization.convert(model_fp32_prepared)
我们来吃一个完整的例子:
# -*- coding:utf-8 -*- # Author:凌逆战 | Never # Date: 2022/10/17 """ 权重和激活都会被量化 """ import torch from torch import nn # 定义一个浮点模型,其中一些层可以被静态量化 class F32Model(torch.nn.Module): def __init__(self): super(F32Model, self).__init__() self.quant = torch.quantization.QuantStub() # QuantStub: 转换张量从浮点到量化 self.conv = nn.Conv2d(1, 1, 1) self.fc = nn.Linear(2, 2, bias=False) self.relu = nn.ReLU() self.dequant = torch.quantization.DeQuantStub() # DeQuantStub: 将量化张量转换为浮点 def forward(self, x): x = self.quant(x) # 手动指定张量: 从浮点转换为量化 x = self.conv(x) x = self.fc(x) x = self.relu(x) x = self.dequant(x) # 手动指定张量: 从量化转换到浮点 return x model_fp32 = F32Model() model_fp32.eval() # 模型必须设置为eval模式,静态量化逻辑才能工作 # 1、如果要部署在ARM上;果要部署在x86 server上 ‘fbgemm’ model_fp32.qconfig = torch.quantization.get_default_qconfig('qnnpack') # 2、在适用的情况下,将一些层进行融合,可以加速 # 常见的融合包括在:DEFAULT_OP_LIST_TO_FUSER_METHOD model_fp32_fused = torch.quantization.fuse_modules(model_fp32, [['fc', 'relu']]) # 3、准备模型,插入observers,观察 activation 和 weight model_fp32_prepared = torch.quantization.prepare(model_fp32_fused) # 4、代表性数据集,获取数据的分布特点,来更好的计算activation的 scale 和 zp input_fp32 = torch.randn(1, 1, 2, 2) # (batch_size, channel, W, H) model_fp32_prepared(input_fp32) # 5、量化模型 model_int8 = torch.quantization.convert(model_fp32_prepared) # 运行模型,相关计算将在int8中进行 output_fp32 = model_fp32(input_fp32) output_int8 = model_int8(input_fp32) print(output_fp32) # tensor([[[[0.6315, 0.0000], # [0.2466, 0.0000]]]], grad_fn=<ReluBackward0>) print(output_int8) # tensor([[[[0.3886, 0.0000], # [0.2475, 0.0000]]]])
# Static quantization of a model consists of the following steps: # Fuse modules # Insert Quant/DeQuant Stubs # Prepare the fused module (insert observers before and after layers) # Calibrate the prepared module (pass it representative data) # Convert the calibrated module (replace with quantized version) import torch from torch import nn import copy backend = "fbgemm" # running on a x86 CPU. Use "qnnpack" if running on ARM. model = nn.Sequential( nn.Conv2d(2,64,3), nn.ReLU(), nn.Conv2d(64, 128, 3), nn.ReLU() ) ## EAGER MODE m = copy.deepcopy(model) m.eval() """Fuse - Inplace fusion replaces the first module in the sequence with the fused module, and the rest with identity modules """ torch.quantization.fuse_modules(m, ['0','1'], inplace=True) # fuse first Conv-ReLU pair torch.quantization.fuse_modules(m, ['2','3'], inplace=True) # fuse second Conv-ReLU pair """Insert stubs""" m = nn.Sequential(torch.quantization.QuantStub(), *m, torch.quantization.DeQuantStub()) """Prepare""" m.qconfig = torch.quantization.get_default_qconfig(backend) torch.quantization.prepare(m, inplace=True) """Calibrate - This example uses random data for convenience. Use representative (validation) data instead. """ with torch.inference_mode(): for _ in range(10): x = torch.rand(1,2, 28, 28) m(x) """Convert""" torch.quantization.convert(m, inplace=True) """Check""" print(m[[1]].weight().element_size()) # 1 byte instead of 4 bytes for FP32 ## FX GRAPH from torch.quantization import quantize_fx m = copy.deepcopy(model) m.eval() qconfig_dict = {"": torch.quantization.get_default_qconfig(backend)} # Prepare model_prepared = quantize_fx.prepare_fx(m, qconfig_dict) # Calibrate - Use representative (validation) data. with torch.inference_mode(): for _ in range(10): x = torch.rand(1,2,28, 28) model_prepared(x) # quantize model_quantized = quantize_fx.convert_fx(model_prepared)
Quantization Aware Training (边训练边量化)
这一部分我用不着,等我需要使用的时候再来补充
保存和加载量化模型
我们先把模型量化
import torch from torch import nn class M(torch.nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(5, 5,bias=True) self.gru = nn.GRU(input_size=5,hidden_size=5,bias=True,) self.relu = nn.ReLU() def forward(self, x): x = self.linear(x) x = self.gru(x) x = self.relu(x) return x m = M().eval() model_int8 = torch.quantization.quantize_dynamic( model=m, # 原始模型 qconfig_spec={nn.Linear, nn.GRU}, # 要动态量化的NN算子 dtype=torch.qint8, inplace=True) # 将权重量化为:float16 \ qint8+
保存/加载量化模型 state_dict
torch.save(model_int8.state_dict(), "./state_dict.pth") model_int8.load_state_dict(torch.load("./state_dict.pth")) print(model_int8)
保存/加载脚本化量化模型 torch.jit.save 和 torch.jit.load
traced_model = torch.jit.trace(model_int8, torch.rand(5, 5)) torch.jit.save(traced_model, "./traced_quant.pt") quantized_model = torch.jit.load("./traced_quant.pt") print(quantized_model)
获取量化模型的参数
其实pytorch获取量化后的模型参数是比较困难的,我们还是以上面的量化模型为例来取参数的值
print(model_int8) # M( # (linear): DynamicQuantizedLinear(in_features=5, out_features=5, dtype=torch.qint8, qscheme=torch.per_tensor_affine) # (gru): DynamicQuantizedGRU(5, 5) # (relu): ReLU() # ) print(model_int8.linear) print(model_int8.gru) print(model_int8.relu)
我们来尝试一下获取线性层的权重和偏置
# print(dir(model_int8.linear)) # 获得对象的所有属性和方法 print(model_int8.linear.weight().int_repr()) # tensor([[ 104, 127, 70, -94, 121], # [ 98, 53, 124, 74, 38], # [-103, -112, 38, 117, 64], # [ -46, -36, 115, 82, -75], # [ -14, -94, 42, -25, 41]], dtype=torch.int8) print(model_int8.linear.bias()) # tensor([ 0.2437, 0.2956, 0.4010, -0.2818, 0.0950], requires_grad=True)
O My God,偏置居然还是浮点类型的,只有权重被量化为了整型。
好的,我们再来获取GRU的权重和偏置
print(dir(model_int8.gru)) print(model_int8.gru.get_weight()["weight_ih_l0"].int_repr()) # int8 print(model_int8.gru.get_weight()["weight_hh_l0"].int_repr()) #int8 print(model_int8.gru.get_bias()["bias_ih_l0"]) # float print(model_int8.gru.get_bias()["bias_hh_l0"]) # float
第一,量化后的取值非常麻烦
第二,静态量化不支持GRU就算了,动态量化偏置还不给我量化了,哎,pytorch的量化真的是还有很长的路要走呀!
参考
【pytorch官方】Quantization(需要非常细心且耐心的去读)
【pytorch官方】Quantization API
【pytorch官方】PyTorch 中的实用量化
【知乎】PyTorch的量化
【CSDN】Pytorch 1.10.2 下模型量化踩坑