PyTorch 中的实用量化

作者：Suraj Subramanian, Mark Saroufim, Jerry Zhang

量化是一种廉价且简单的方法，可以使您的 DNN 运行得更快，并降低内存需求。PyTorch 提供了几种不同的模型量化方法。在这篇博文中，我们将快速奠定深度学习中量化的基础，然后看看每种技术在实践中是如何应用的。最后，我们将结合文献中的建议，为您在工作流程中使用量化提供建议。

图 1. PyTorch <3 量化

量化基础
在 PyTorch 中
敏感性分析
工作流程建议
- 注意事项
参考文献

量化基础

如果有人问你几点钟，你不会回答“10:14:34:430705”，而是可能说“十点一刻”。

量化根源于信息压缩；在深度网络中，它指的是降低其权重和/或激活值的数值精度。

过参数化的 DNN 具有更多自由度，这使其成为信息压缩的良好候选 [1]。量化模型后，通常会发生两件事——模型变得更小，运行效率更高。硬件供应商明确支持更快的 8 位数据处理（相比 32 位数据），从而提高吞吐量。更小的模型具有更低的内存占用和功耗 [2]，这对于边缘部署至关重要。

映射函数

映射函数正如您所猜测的——一个将值从浮点空间映射到整数空间的函数。常用的映射函数是由 $Q(r) = round(r/S + Z)$ 给出的线性变换，其中 $r$ 是输入， $S, Z$ 是量化参数。

要重新转换回浮点空间，逆函数由 $\tilde r = (Q(r) - Z) \cdot S$ 给出。

$\tilde r \neq r$ ，它们的差值构成了量化误差。

量化参数

映射函数由缩放因子 $S$ 和零点 $Z$ 参数化。

$S$ 简单地是输入范围与输出范围的比率 $S = \frac{\beta - \alpha}{\beta_q - \alpha_q}$

其中 [ $\alpha, \beta$ ] 是输入的截取范围，即允许输入的边界。[ $\alpha_q, \beta_q$ ] 是量化输出空间中映射到的范围。对于 8 位量化，输出范围 $\beta_q - \alpha_q <= (2^8 - 1)$ 。

$Z$ 作为偏置，确保输入空间中的 0 完美映射到量化空间中的 0。 $Z = -(\frac{\alpha}{S} - \alpha_q)$

校准

选择输入截取范围的过程称为校准。最简单的技术（也是 PyTorch 中的默认技术）是记录运行中的最小值和最大值并将它们赋给 $\alpha$ 和 $\beta$ 。TensorRT 还使用熵最小化（KL 散度）、均方误差最小化或输入范围的百分位数。

在 PyTorch 中，Observer 模块（代码）收集输入值的统计信息并计算 qparams $S, Z$ 。不同的校准方案会导致不同的量化输出，最好通过实验验证哪种方案最适合您的应用和架构（稍后会详细介绍）。

from torch.quantization.observer import MinMaxObserver, MovingAverageMinMaxObserver, HistogramObserver
C, L = 3, 4
normal = torch.distributions.normal.Normal(0,1)
inputs = [normal.sample((C, L)), normal.sample((C, L))]
print(inputs)

# >>>>>
# [tensor([[-0.0590,  1.1674,  0.7119, -1.1270],
#          [-1.3974,  0.5077, -0.5601,  0.0683],
#          [-0.0929,  0.9473,  0.7159, -0.4574]]]),

# tensor([[-0.0236, -0.7599,  1.0290,  0.8914],
#          [-1.1727, -1.2556, -0.2271,  0.9568],
#          [-0.2500,  1.4579,  1.4707,  0.4043]])]

observers = [MinMaxObserver(), MovingAverageMinMaxObserver(), HistogramObserver()]
for obs in observers:
  for x in inputs: obs(x) 
  print(obs.__class__.__name__, obs.calculate_qparams())

# >>>>>
# MinMaxObserver (tensor([0.0112]), tensor([124], dtype=torch.int32))
# MovingAverageMinMaxObserver (tensor([0.0101]), tensor([139], dtype=torch.int32))
# HistogramObserver (tensor([0.0100]), tensor([106], dtype=torch.int32))

仿射和对称量化方案

仿射或非对称量化方案将输入范围指定为观察到的最小值和最大值。仿射方案通常提供更紧凑的截取范围，对于量化非负激活值非常有用（如果您的输入张量从不包含负值，则不需要输入范围包含负值）。范围计算为 $\alpha = min(r), \beta = max(r)$ 。将仿射量化用于权重张量时，会导致推理计算开销更大 [3]。

对称量化方案将输入范围中心设为 0，从而无需计算零点偏移。范围计算为 $-\alpha = \beta = max(|max(r)|,|min(r)|)$ 。对于偏斜信号（如非负激活值），这可能导致较差的量化分辨率，因为截取范围包含了从未出现在输入中的值（参见下面的 pyplot）。

act =  torch.distributions.pareto.Pareto(1, 10).sample((1,1024))
weights = torch.distributions.normal.Normal(0, 0.12).sample((3, 64, 7, 7)).flatten()

def get_symmetric_range(x):
  beta = torch.max(x.max(), x.min().abs())
  return -beta.item(), beta.item()

def get_affine_range(x):
  return x.min().item(), x.max().item()

def plot(plt, data, scheme):
  boundaries = get_affine_range(data) if scheme == 'affine' else get_symmetric_range(data)
  a, _, _ = plt.hist(data, density=True, bins=100)
  ymin, ymax = np.quantile(a[a>0], [0.25, 0.95])
  plt.vlines(x=boundaries, ls='--', colors='purple', ymin=ymin, ymax=ymax)

fig, axs = plt.subplots(2,2)
plot(axs[0, 0], act, 'affine')
axs[0, 0].set_title("Activation, Affine-Quantized")

plot(axs[0, 1], act, 'symmetric')
axs[0, 1].set_title("Activation, Symmetric-Quantized")

plot(axs[1, 0], weights, 'affine')
axs[1, 0].set_title("Weights, Affine-Quantized")

plot(axs[1, 1], weights, 'symmetric')
axs[1, 1].set_title("Weights, Symmetric-Quantized")
plt.show()

图 2. 仿射和对称方案的截取范围（紫色部分）

在 PyTorch 中，您可以在初始化 Observer 时指定仿射或对称方案。请注意，并非所有 Observer 都支持这两种方案。

for qscheme in [torch.per_tensor_affine, torch.per_tensor_symmetric]:
  obs = MovingAverageMinMaxObserver(qscheme=qscheme)
  for x in inputs: obs(x)
  print(f"Qscheme: {qscheme} | {obs.calculate_qparams()}")

# >>>>>
# Qscheme: torch.per_tensor_affine | (tensor([0.0101]), tensor([139], dtype=torch.int32))
# Qscheme: torch.per_tensor_symmetric | (tensor([0.0109]), tensor([128]))

Per-Tensor 和 Per-Channel 量化方案

量化参数可以针对层的整个权重张量作为一个整体计算，也可以分别为每个通道计算。在 per-tensor 方案中，相同的截取范围应用于层中的所有通道

图 3. Per-Channel 为每个通道使用一组 qparams。Per-Tensor 为整个张量使用相同的 qparams。

对于权重量化，symmetric-per-channel 量化提供了更好的精度；per-tensor 量化性能较差，这可能是由于 batchnorm 折叠导致不同通道的 conv 权重方差较大 [3]。

from torch.quantization.observer import MovingAveragePerChannelMinMaxObserver
obs = MovingAveragePerChannelMinMaxObserver(ch_axis=0)  # calculate qparams for all `C` channels separately
for x in inputs: obs(x)
print(obs.calculate_qparams())

# >>>>>
# (tensor([0.0090, 0.0075, 0.0055]), tensor([125, 187,  82], dtype=torch.int32))

后端引擎

目前，量化算子通过 FBGEMM 后端在 x86 机器上运行，或在 ARM 机器上使用 QNNPACK 原语。服务器 GPU 的后端支持（通过 TensorRT 和 cuDNN）即将推出。了解有关将量化扩展到自定义后端的更多信息：RFC-0019。

backend = 'fbgemm' if x86 else 'qnnpack'
qconfig = torch.quantization.get_default_qconfig(backend)  
torch.backends.quantized.engine = backend

QConfig

QConfig (代码) NamedTuple 存储用于量化激活值和权重的 Observer 和量化方案。

请务必传递 Observer 类（而不是实例），或一个可以返回 Observer 实例的可调用对象。使用 with_args() 来覆盖默认参数。

my_qconfig = torch.quantization.QConfig(
  activation=MovingAverageMinMaxObserver.with_args(qscheme=torch.per_tensor_affine),
  weight=MovingAveragePerChannelMinMaxObserver.with_args(qscheme=torch.qint8)
)
# >>>>>
# QConfig(activation=functools.partial(<class 'torch.ao.quantization.observer.MovingAverageMinMaxObserver'>, qscheme=torch.per_tensor_affine){}, weight=functools.partial(<class 'torch.ao.quantization.observer.MovingAveragePerChannelMinMaxObserver'>, qscheme=torch.qint8){})

在 PyTorch 中

PyTorch 允许您通过几种不同的方式量化模型，具体取决于

您是偏好灵活但手动的方式，还是受限的自动处理方式（Eager Mode 对比 FX Graph Mode）
量化激活值（层输出）的 qparams 是预先计算好的供所有输入使用，还是针对每个输入重新计算（静态对比动态），
qparams 是在有或无再训练的情况下计算的（量化感知训练 对比 训练后量化）

FX Graph Mode 自动融合符合条件的模块，插入 Quant/DeQuant 桩，校准模型并返回量化模块——所有这些都只需两次方法调用——但仅适用于可符号跟踪的网络。下面的示例包含了使用 Eager Mode 和 FX Graph Mode 的调用，供比较参考。

在 DNN 中，适合量化的候选对象是 FP32 权重（层参数）和激活值（层输出）。量化权重可以减小模型大小。量化激活值通常会加快推理速度。

例如，50 层的 ResNet 网络拥有约 2600 万个权重参数，并在前向传播中计算约 1600 万个激活值。

训练后动态/仅权重量化

在这里，模型的权重是预先量化的；激活值在推理过程中即时量化（“动态”）。这是所有方法中最简单的一种，只需一行 API 调用 torch.quantization.quantize_dynamic。目前动态量化仅支持 Linear 和 Recurrent（LSTM、GRU、RNN）层。

(+) 可以获得更高的精度，因为每个输入的截取范围都经过精确校准 [1]。

(+) 对于 LSTM 和 Transformer 等模型，动态量化是首选，因为从内存写入/读取模型权重会占据主要带宽 [4]。

(-) 在运行时校准和量化每一层的激活值会增加计算开销。

import torch
from torch import nn

# toy model
m = nn.Sequential(
  nn.Conv2d(2, 64, (8,)),
  nn.ReLU(),
  nn.Linear(16,10),
  nn.LSTM(10, 10))

m.eval()

## EAGER MODE
from torch.quantization import quantize_dynamic
model_quantized = quantize_dynamic(
    model=m, qconfig_spec={nn.LSTM, nn.Linear}, dtype=torch.qint8, inplace=False
)

## FX MODE
from torch.quantization import quantize_fx
qconfig_dict = {"": torch.quantization.default_dynamic_qconfig}  # An empty key denotes the default applied to all modules
model_prepared = quantize_fx.prepare_fx(m, qconfig_dict)
model_quantized = quantize_fx.convert_fx(model_prepared)

训练后静态量化 (PTQ)

PTQ 也预先量化模型权重，但不是即时校准激活值，而是使用验证数据预先校准并固定（“静态”）截取范围。在推理过程中，激活值在操作之间保持量化精度。大约 100 个代表性数据的 mini-batch 足以校准 observer [2]。下面的示例为方便起见在校准中使用了随机数据——在您的应用中使用随机数据会导致 qparams 不佳。

PTQ flowchart
图 4. 训练后静态量化的步骤

模块融合将多个连续模块（例如：[Conv2d, BatchNorm, ReLU]）合并为一个。融合模块意味着编译器只需要运行一个内核而不是多个；这可以加快速度并通过减少量化误差来提高精度。

(+) 静态量化比动态量化推理更快，因为它消除了层之间的浮点<->整数转换开销。

(-) 静态量化模型可能需要定期重新校准以应对分布漂移。

# Static quantization of a model consists of the following steps:

#     Fuse modules
#     Insert Quant/DeQuant Stubs
#     Prepare the fused module (insert observers before and after layers)
#     Calibrate the prepared module (pass it representative data)
#     Convert the calibrated module (replace with quantized version)

import torch
from torch import nn
import copy

backend = "fbgemm"  # running on a x86 CPU. Use "qnnpack" if running on ARM.

model = nn.Sequential(
     nn.Conv2d(2,64,3),
     nn.ReLU(),
     nn.Conv2d(64, 128, 3),
     nn.ReLU()
)

## EAGER MODE
m = copy.deepcopy(model)
m.eval()
"""Fuse
- Inplace fusion replaces the first module in the sequence with the fused module, and the rest with identity modules
"""
torch.quantization.fuse_modules(m, ['0','1'], inplace=True) # fuse first Conv-ReLU pair
torch.quantization.fuse_modules(m, ['2','3'], inplace=True) # fuse second Conv-ReLU pair

"""Insert stubs"""
m = nn.Sequential(torch.quantization.QuantStub(), 
                  *m, 
                  torch.quantization.DeQuantStub())

"""Prepare"""
m.qconfig = torch.quantization.get_default_qconfig(backend)
torch.quantization.prepare(m, inplace=True)

"""Calibrate
- This example uses random data for convenience. Use representative (validation) data instead.
"""
with torch.inference_mode():
  for _ in range(10):
    x = torch.rand(1,2, 28, 28)
    m(x)
    
"""Convert"""
torch.quantization.convert(m, inplace=True)

"""Check"""
print(m[[1]].weight().element_size()) # 1 byte instead of 4 bytes for FP32


## FX GRAPH
from torch.quantization import quantize_fx
m = copy.deepcopy(model)
m.eval()
qconfig_dict = {"": torch.quantization.get_default_qconfig(backend)}
# Prepare
model_prepared = quantize_fx.prepare_fx(m, qconfig_dict)
# Calibrate - Use representative (validation) data.
with torch.inference_mode():
  for _ in range(10):
    x = torch.rand(1,2,28, 28)
    model_prepared(x)
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

量化感知训练 (QAT)

QAT flowchart
图 5. 量化感知训练的步骤

PTQ 方法非常适用于大型模型，但在小型模型中精度会受到影响 [[6]]。这当然是由于将模型从 FP32 适配到 INT8 领域时数值精度损失所致 (图 6(a))。QAT 通过将此量化误差包含在训练损失中来解决此问题，从而训练一个 INT8 优先的模型。

Fig. 6: Comparison of PTQ and QAT
图 6. PTQ 和 QAT 收敛比较 [3]

所有权重和偏置都存储在 FP32 中，反向传播照常进行。然而在前向传播中，量化通过 FakeQuantize 模块进行内部模拟。它们被称为“假”是因为它们对数据进行量化后立即反量化，添加了类似于量化推理中可能遇到的量化噪声。最终损失因此考虑了任何预期的量化误差。在此基础上进行优化，模型可以在损失函数中识别更宽泛的区域 (图 6(b))，并识别出 FP32 参数，使其量化为 INT8 后不会显著影响精度。

Fake Quantization in the forward and backward pass
图 7. 前向和反向传播中的假量化
图片来源: https://developer.nvidia.com/blog/achieving-fp32-accuracy-for-int8-inference-using-quantization-aware-training-with-tensorrt

(+) QAT 获得的精度高于 PTQ。

(+) 可以在模型训练期间学习 Qparams，以获得更精细的精度（参见 LearnableFakeQuantize）

(-) 在 QAT 中再训练模型的计算开销可能高达数百个 epoch [1]

# QAT follows the same steps as PTQ, with the exception of the training loop before you actually convert the model to its quantized version

import torch
from torch import nn

backend = "fbgemm"  # running on a x86 CPU. Use "qnnpack" if running on ARM.

m = nn.Sequential(
     nn.Conv2d(2,64,8),
     nn.ReLU(),
     nn.Conv2d(64, 128, 8),
     nn.ReLU()
)

"""Fuse"""
torch.quantization.fuse_modules(m, ['0','1'], inplace=True) # fuse first Conv-ReLU pair
torch.quantization.fuse_modules(m, ['2','3'], inplace=True) # fuse second Conv-ReLU pair

"""Insert stubs"""
m = nn.Sequential(torch.quantization.QuantStub(), 
                  *m, 
                  torch.quantization.DeQuantStub())

"""Prepare"""
m.train()
m.qconfig = torch.quantization.get_default_qconfig(backend)
torch.quantization.prepare_qat(m, inplace=True)

"""Training Loop"""
n_epochs = 10
opt = torch.optim.SGD(m.parameters(), lr=0.1)
loss_fn = lambda out, tgt: torch.pow(tgt-out, 2).mean()
for epoch in range(n_epochs):
  x = torch.rand(10,2,24,24)
  out = m(x)
  loss = loss_fn(out, torch.rand_like(out))
  opt.zero_grad()
  loss.backward()
  opt.step()

"""Convert"""
m.eval()
torch.quantization.convert(m, inplace=True)

敏感性分析

并非所有层对量化的响应都相同，有些层比其他层对精度下降更敏感。识别最小化精度下降的最佳层组合非常耗时，因此 [3] 建议进行一次一层敏感性分析，以识别哪些层最敏感，并在这些层上保留 FP32 精度。在他们的实验中，仅跳过 2 个 conv 层（MobileNet v1 中总共有 28 个）即可获得接近 FP32 的精度。使用 FX Graph Mode，我们可以轻松创建自定义 qconfigs 来实现这一点

# ONE-AT-A-TIME SENSITIVITY ANALYSIS 

for quantized_layer, _ in model.named_modules():
  print("Only quantizing layer: ", quantized_layer)

  # The module_name key allows module-specific qconfigs. 
  qconfig_dict = {"": None, 
  "module_name":[(quantized_layer, torch.quantization.get_default_qconfig(backend))]}

  model_prepared = quantize_fx.prepare_fx(model, qconfig_dict)
  # calibrate
  model_quantized = quantize_fx.convert_fx(model_prepared)
  # evaluate(model)

另一种方法是比较 FP32 和 INT8 层的统计信息；常用的衡量指标是 SQNR（信号量化噪声比）和均方误差。这种比较分析也有助于指导进一步的优化。

Fig 8. Comparing model weights and activations
图 8. 模型权重和激活值比较

PyTorch 在 Numeric Suite 下提供了有助于进行此分析的工具。从完整教程了解更多关于使用 Numeric Suite 的信息。

# extract from https://pytorch.ac.cn/tutorials/prototype/numeric_suite_tutorial.html
import torch.quantization._numeric_suite as ns

def SQNR(x, y):
    # Higher is better
    Ps = torch.norm(x)
    Pn = torch.norm(x-y)
    return 20*torch.log10(Ps/Pn)

wt_compare_dict = ns.compare_weights(fp32_model.state_dict(), int8_model.state_dict())
for key in wt_compare_dict:
    print(key, compute_error(wt_compare_dict[key]['float'], wt_compare_dict[key]['quantized'].dequantize()))

act_compare_dict = ns.compare_model_outputs(fp32_model, int8_model, input_data)
for key in act_compare_dict:
    print(key, compute_error(act_compare_dict[key]['float'][0], act_compare_dict[key]['quantized'][0].dequantize()))

工作流程建议

Suggested quantization workflow
图 9. 建议的量化工作流程

点击查看大图

注意事项

大型模型（参数量大于 1000 万）对量化误差更鲁棒。[2]
从 FP32 检查点量化模型比从头开始训练 INT8 模型可以提供更好的精度。[2]
对模型运行时进行性能分析是可选的，但它可以帮助识别限制推理速度的层。
动态量化是一个简单的第一步，特别是如果您的模型包含许多 Linear 或 Recurrent 层。
使用 symmetric-per-channel 量化以及 MinMax observer 来量化权重。使用 affine-per-tensor 量化以及 MovingAverageMinMax observer 来量化激活值 [2, 3]
使用 SQNR 等指标来识别哪些层最容易受到量化误差的影响。在这些层上关闭量化。
使用 QAT 进行微调，持续时间约为原始训练计划的 10%，学习率调度采用退火策略，起始学习率为初始训练学习率的 1%。[3]
如果上述工作流程对您不起作用，我们希望了解更多信息。请发布一个帖子，详细说明您的代码（模型架构、精度指标、尝试的技术）。请随意抄送我 @suraj.pt。

内容很多，感谢您坚持读完！接下来，我们将看看如何量化使用动态控制结构（if-else、循环）的“真实世界”模型。这些元素使得模型无法进行符号跟踪，从而难以直接对模型进行开箱即用的量化。在本系列的下一篇博文中，我们将亲自动手处理一个包含大量循环和 if-else 块、甚至在 forward 调用中使用了第三方库的模型。

我们还将介绍 PyTorch 量化中一个很酷的新功能，称为 Define-by-Run，它试图通过仅要求模型的计算图的子集不包含动态流来缓解此限制。查看 PTDD’21 上的 Define-by-Run 海报以先睹为快。

参考文献

[1] Gholami, A., Kim, S., Dong, Z., Yao, Z., Mahoney, M. W., & Keutzer, K. (2021). A survey of quantization methods for efficient neural network inference. arXiv preprint arXiv:2103.13630。

[2] Krishnamoorthi, R. (2018). Quantizing deep convolutional networks for efficient inference: A whitepaper. arXiv preprint arXiv:1806.08342。

[3] Wu, H., Judd, P., Zhang, X., Isaev, M., & Micikevicius, P. (2020). Integer quantization for deep learning inference: Principles and empirical evaluation. arXiv preprint arXiv:2004.09602。

[4] PyTorch 量化文档