量化是一种廉价且简单的方法，可以使您的深度神经网络（DNN）运行得更快，并降低内存需求。PyTorch 提供了几种不同的模型量化方法。在这篇博文中，我们将快速构建深度学习量化的基础，然后看看每种技术在实践中是如何应用的。最后，我们将总结文献中关于在工作流程中使用量化的建议。
图 1. PyTorch <3 量化
目录
量化基础
如果有人问你现在几点，你不会回答“10:14:34:430705”，你可能会说“10点过一刻”。
量化源于信息压缩；在深度网络中，它指的是降低权重和/或激活值的数值精度。
过参数化的 DNN 具有更多的自由度，这使其成为信息压缩的良好候选对象 [1]。当你量化一个模型时，通常会发生两件事——模型变小了，运行效率更高了。硬件供应商明确支持比 32 位数据更快地处理 8 位数据，从而获得更高的吞吐量。更小的模型具有更低的内存占用和功耗 [2]，这对边缘部署至关重要。
映射函数
映射函数正如你所猜想的那样——一个将值从浮点空间映射到整数空间的函数。一种常用的映射函数是线性变换 ，其中
是输入，
是量化参数。
要转换回浮点空间，逆函数为 。
，它们的差值构成量化误差。
量化参数
映射函数由缩放因子 和零点
进行参数化。
仅仅是输入范围与输出范围的比率
其中 [] 是输入的截断范围，即允许输入的边界。[
] 是它所映射到的量化输出空间范围。对于 8 位量化，输出范围
。
充当偏差，以确保输入空间中的 0 完美映射到量化空间中的 0。
校准
选择输入截断范围的过程称为校准。最简单的技术（也是 PyTorch 中的默认设置）是记录运行时的最小值和最大值，并将它们分配给 和
。TensorRT 还使用熵最小化（KL 散度）、均方误差最小化或输入范围的分位数。
在 PyTorch 中，Observer 模块（代码）收集关于输入值的统计信息并计算量化参数 。不同的校准方案会导致不同的量化输出，最好通过实验验证哪种方案最适合您的应用程序和架构（稍后会详细介绍）。
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))
仿射和对称量化方案
仿射或非对称量化方案将输入范围分配给观测到的最小值和最大值。仿射方案通常提供更紧凑的截断范围，适用于量化非负激活（如果输入张量从不为负，则不需要输入范围包含负值）。范围计算为 。当用于权重张量时，仿射量化会导致计算开销更大的推理 [3]。
对称量化方案将输入范围以 0 为中心，消除了计算零点偏移的需要。范围计算为 。对于偏斜信号（如非负激活），这可能会导致较差的量化分辨率，因为截断范围包含了输入中从不出现的值（见下方的 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 时指定仿射或对称方案。请注意，并非所有观察者都支持这两种方案。
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）方案中，相同的截断范围应用于层中的所有通道。
图 3. 逐通道方案为每个通道使用一组量化参数。逐张量方案为整个张量使用相同的量化参数。
对于权重量化，对称逐通道量化可提供更好的精度；逐张量量化表现不佳，这可能是由于 Batchnorm 折叠后跨通道卷积权重的高方差所致 [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 实例的可调用对象。使用 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 模式 与 FX 图模式），
- 量化激活（层输出）的量化参数是预先为所有输入计算的，还是在每次输入时重新计算（静态 与 动态），
- 量化参数是在有无重训练的情况下计算的（量化感知训练 与 训练后量化）。
FX 图模式会自动融合符合条件的模块，插入量化/反量化存根，校准模型并返回量化后的模块——所有这些只需两个方法调用——但仅适用于可进行符号追踪的网络。下面的示例包含使用 Eager 模式和 FX 图模式的调用以供对比。
在 DNN 中，符合量化条件的候选者是 FP32 权重（层参数）和激活（层输出）。量化权重可减小模型大小。量化激活通常会加快推理速度。
例如，50 层的 ResNet 网络具有约 2600 万个权重参数，并在前向传播中计算约 1600 万个激活。
训练后动态量化/仅权重量化
此处模型的权重是预量化的；激活在推理过程中即时（“动态”）量化。这是所有方法中最简单的一种，只需在 torch.quantization.quantize_dynamic 中进行一次 API 调用。目前动态量化仅支持线性层和循环层（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 个代表性数据的迷你批次足以校准观察者 [2]。下面的示例为了方便使用了随机数据进行校准——在您的应用程序中使用它会导致错误的量化参数。
图 4. 训练后静态量化的步骤
模块融合将多个连续模块（例如：[Conv2d, BatchNorm, ReLU]）组合成一个。融合模块意味着编译器只需要运行一个核而不是多个；这加快了速度，并通过减少量化误差提高了精度。
(+) 静态量化的推理速度比动态量化快，因为它消除了层之间 float<->int 的转换成本。
(-) 静态量化模型可能需要定期重新校准，以应对分布漂移。
# 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)
图 5. 量化感知训练的步骤
PTQ 方法非常适合大型模型，但小型模型的精度会受到影响 [[6]]。这当然是因为将模型从 FP32 调整到 INT8 领域时损失了数值精度 (图 6(a))。QAT 通过将这种量化误差包含在训练损失中来解决这个问题，从而训练一个 INT8 优先的模型。
图 6. PTQ 和 QAT 收敛对比 [3]
所有权重和偏差都以 FP32 存储，反向传播按正常方式进行。然而，在前向传播中，量化是通过 FakeQuantize 模块在内部模拟的。它们之所以被称为“伪”，是因为它们量化并立即反量化数据，增加了类似于量化推理期间可能遇到的量化噪声。因此，最终损失考虑了任何预期的量化误差。在此基础上进行优化，使模型能够在损失函数中识别出更广泛的区域 (图 6(b))，并识别 FP32 参数，使得将它们量化为 INT8 不会显著影响精度。
图 7. 前向和反向传播中的伪量化
图片来源：https://developer.nvidia.com/blog/achieving-fp32-accuracy-for-int8-inference-using-quantization-aware-training-with-tensorrt
(+) QAT 产生比 PTQ 更高的精度。
(+) 量化参数可以在模型训练期间学习，以获得更精细的精度（参见 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 个卷积层（MobileNet v1 总共 28 层中）就使他们获得了接近 FP32 的精度。使用 FX 图模式，我们可以轻松创建自定义 qconfig 来实现这一点。
# 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（信号量化噪声比）和均方误差。此类比较分析也有助于指导进一步的优化。
图 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()))
工作流程建议
图 9. 建议的量化工作流程
注意事项
- 大型（10M+ 参数）模型对量化误差更稳健。 [2]
- 从 FP32 检查点量化模型比从头开始训练 INT8 模型能提供更好的精度。[2]
- 分析模型运行时间是可选的，但它可以帮助识别推理瓶颈层。
- 动态量化是简单的第一步，特别是如果您的模型有很多线性或循环层。
- 对于量化权重，使用带有
MinMax观察者的对称逐通道量化。对于量化激活，使用带有MovingAverageMinMax观察者的仿射逐张量量化 [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 量化文档