1. 首頁 > 生活百科 > >

Pytorch模型量化( 二 )

生活百科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)輸出
    Pytorch模型量化

    文章插圖
    Pytorch模型量化

    文章插圖
    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)

    推薦閱讀