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中的值

　　为了实现量化，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中，选择合适的scale和zp的工作就由各种observer来完成。

Tensor的量化支持两种模式：per tensor 和 per channel。

• Per tensor：是说一个tensor里的所有value按照同一种方式去scale和offset；
• Per channel：是对于tensor的某一个维度（通常是channel的维度）上的值按照一种方式去scale和offset，也就是一个tensor里有多种不同的scale和offset的方式（组成一个vector），如此以来，在量化的时候相比per tensor的方式会引入更少的错误。PyTorch目前支持conv2d()、conv3d()、linear()的per channel量化。

在我们正式了解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]]]])