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

在我们正式了解pytorch模型量化前我们再来检查一下pytorch的官方量化是否能满足我们的需求，如果不能，后面的都不需要看了
静态量化动态量化nn.linearYYnn.Conv1d/2d/3dYN (因为pytorch认为卷积参数来了个太小了，对卷积核进行量化会造成更多损失，所以pytorch选择不量化)nn.LSTMN(LSTM的好像又可以了，官方给出了一个例子，传送门)Ynn.GRUNYnn.RNNCellNYnn.GRUCellNYnn.LSTMCellNYnn.EmbeddingBagY(激活在fp32)Ynn.EmbeddingYNnn.MultiheadAttentionNNActivations大部分支持不变，计算停留在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 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)输出