V2AS问路
Pytorch模型量化
阅读原文时间：2023年07月08日阅读：5

在深度学习中，量化指的是使用更少的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的量化。

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=)
tensor([[[[-0.5371, 0.0713],
[-0.2040, 0.4126]]]])

Post Training Static Quantization (训练后静态量化)

　　静态量化需要把模型的权重和激活都进行量化，静态量化需要把训练集或者和训练集分布类似的数据喂给模型(注意没有反向传播)，然后通过每个op输入的分布 来计算activation的量化参数（scale和zp）——称之为Calibrate（定标），因为静态量化的前向推理过程自始至终都是int计算，activation需要确保一个op的输入符合下一个op的输入。

PyTorch会使用以下5步来完成模型的静态量化：

合并一些可以合并的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)

一旦合并成功，那么原始网络中的fc就会被替换为新的合并后的module（因为其是list中的第一个元素），而relu（list中剩余的元素）会被替换为nn.Identity()，这个模块是个占位符，直接输出输入。举个例子，对于下面的一个小网络：

import torch
from torch import nn

class 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 x

model_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

qconfig要设置到模型或者Module上。

#如果要部署在x86 server上
model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm')

#如果要部署在ARM上
model_fp32.qconfig = torch.quantization.get_default_qconfig('qnnpack')

x86和arm之外目前不支持。

prepare用来给每个子module插入Observer，用来收集和定标数据。

以activation的observer为例，观察输入数据得到 四元组中的 min_val 和 max_val，至少观察个几百个迭代的数据吧，然后由这四元组得到 scale 和 zp 这两个参数的值。

model_fp32_prepared= torch.quantization.prepare(model_fp32_fused)

这一步不是训练。是为了获取数据的分布特点，来更好的计算activation的 scale 和 zp 。至少要喂上几百个迭代的数据。

#至少观察个几百迭代
for data in data_loader:
model_fp32_prepared(data)

第四步完成后，各个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 torch
from 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 x

model_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_METHOD

model_fp32_fused = torch.quantization.fuse_modules(model_fp32, [['fc', 'relu']])

3、准备模型，插入observers，观察 activation 和 weight

model_fp32_prepared = torch.quantization.prepare(model_fp32_fused)

4、代表性数据集，获取数据的分布特点，来更好的计算activation的 scale 和 zp

input_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=)

print(output_int8)

tensor([[[[0.3886, 0.0000],

[0.2475, 0.0000]]]])

Quantization Aware Training (边训练边量化)

这一部分我用不着，等我需要使用的时候再来补充

保存和加载量化模型

我们先把模型量化

import torch
from torch import nn

class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(5, 5,bias=True)
self.gru = nn.GRU(input_size=5,hidden_size=5,bias=True,)
self.relu = nn.ReLU()

def forward(self, x):  
    x = self.linear(x)  
    x = self.gru(x)  
    x = self.relu(x)  
    return x

m = M().eval()
model_int8 = torch.quantization.quantize_dynamic(
model=m, # 原始模型
qconfig_spec={nn.Linear,
nn.GRU}, # 要动态量化的NN算子
dtype=torch.qint8, inplace=True) # 将权重量化为：float16 \ qint8+

torch.save(model_int8.state_dict(), "./state_dict.pth")
model_int8.load_state_dict(torch.load("./state_dict.pth"))
print(model_int8)

traced_model = torch.jit.trace(model_int8, torch.rand(5, 5))
torch.jit.save(traced_model, "./traced_quant.pt")
quantized_model = torch.jit.load("./traced_quant.pt")
print(quantized_model)

获取量化模型的参数

其实pytorch获取量化后的模型参数是比较困难的，我们还是以上面的量化模型为例来取参数的值

print(model_int8)

M(

(linear): DynamicQuantizedLinear(in_features=5, out_features=5, dtype=torch.qint8, qscheme=torch.per_tensor_affine)

(gru): DynamicQuantizedGRU(5, 5)

(relu): ReLU()

)

print(model_int8.linear)
print(model_int8.gru)
print(model_int8.relu)

我们来尝试一下获取线性层的权重和偏置

# print(dir(model_int8.linear))　　# 获得对象的所有属性和方法
print(model_int8.linear.weight().int_repr())

tensor([[ 104, 127, 70, -94, 121],

[ 98, 53, 124, 74, 38],

[-103, -112, 38, 117, 64],

[ -46, -36, 115, 82, -75],

[ -14, -94, 42, -25, 41]], dtype=torch.int8)

print(model_int8.linear.bias())

tensor([ 0.2437, 0.2956, 0.4010, -0.2818, 0.0950], requires_grad=True)

O My God，偏置居然还是浮点类型的，只有权重被量化为了整型。

好的，我们再来获取GRU的权重和偏置

print(dir(model_int8.gru))
print(model_int8.gru.get_weight()["weight_ih_l0"].int_repr()) # int8
print(model_int8.gru.get_weight()["weight_hh_l0"].int_repr()) #int8
print(model_int8.gru.get_bias()["bias_ih_l0"]) # float
print(model_int8.gru.get_bias()["bias_hh_l0"]) # float

第一，别问我别问我为什么取值这么麻烦，你以为我想？？？

第二，静态量化不支持GRU就算了，动态量化偏置还不给我量化了，哎，pytorch的量化真的是还有很长的路要走呀！

参考

【pytorch官方】Quantization（需要非常细心且耐心的去读）

【pytorch官方】Quantization API

【知乎】PyTorch的量化

手机扫一扫

移动阅读更方便

阿里云服务器
腾讯云服务器
七牛云服务器

你可能感兴趣的文章

使用Triton部署chatglm2-6b模型
重磅｜ Falcon 180B 正式在 Hugging Face Hub 上发布！
使用 Transformers 为多语种语音识别任务微调 Whisper 模型

4

基于bert-base-chinese训练bert模型（最后附上整体代码）

5

TheTransformerPlaybookforNLPandLanguageUnderstanding

6

带你上手基于Pytorch和Transformers的中文NLP训练框架

7

torch.nn基础学习教程 | PyTorch nn Basic Tutorial

8

【pytorch】从零开始，利用yolov5、crnn+ctc进行车牌识别

9

循环神经网络RNN完全解析：从基础理论到PyTorch实战