1. 首頁 > 生活百科 > >

Pytorch模型量化( 三 )

生活百科Pytorch模型量化

一旦合并成功，那么原始網絡中的fc就會被替換為新的合并后的module（因為其是list中的第一個元素），而relu（list中剩余的元素）會被替換為nn.Identity()，這個模塊是個占位符，直接輸出輸入 。舉個例子，對于下面的一個小網絡：
【Pytorch模型量化】import torchfrom torch import nnclass 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 xmodel_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、設置qconfigqconfig要設置到模型或者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、prepareprepare用來給每個子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 torchfrom 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 xmodel_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_METHODmodel_fp32_fused = torch.quantization.fuse_modules(model_fp32, [['fc', 'relu']])# 3、準備模型，插入observers ， 觀察 activation 和 weightmodel_fp32_prepared = torch.quantization.prepare(model_fp32_fused)# 4、代表性數據集，獲取數據的分布特點，來更好的計算activation的 scale 和 zpinput_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]]]])

推薦閱讀