Pytorch模型量化( 二 ) _生活百科

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 torchfrom torch import nnclass 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 xif __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]]]])Post Training Static Quantization (训练后静态量化)静态量化需要把模型的权重和激活都进行量化，静态量化需要把训练集或者和训练集分布类似的数据喂给模型(注意没有反向传播)，然后通过每个op输入的分布来计算activation的量化参数（scale和zp）——称之为Calibrate（定标），因为静态量化的前向推理过程自始至终都是int计算，activation需要确保一个op的输入符合下一个op的输入。
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传递下面的参数就会合并网络中的conv1、bn1、relu1：
torch.quantization.fuse_modules(F32Model, [['fc', 'relu']], inplace=True)