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的量化 。
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 torchx = 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 tensorxq = 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 tensorxdq = 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.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的参数量很大,收益更高) 。

推荐阅读