量化是一种经济且简便的方法，可以使您的 DNN 运行更快，并降低内存要求。PyTorch 提供了几种不同的模型量化方法。在这篇博文中，我们将（快速）介绍深度学习中量化的基础知识，然后看看每种技术在实践中的应用。最后，我们将结合文献中的建议，为您的工作流程提供量化建议。
图 1. PyTorch <3 量化
目录
量化基础
如果有人问您现在几点，您不会回答“10:14:34:430705”，而是可能会说“十点一刻”。
量化源于信息压缩；在深度网络中，它指的是降低其权重和/或激活的数值精度。
过度参数化的 DNN 具有更多的自由度，这使得它们成为信息压缩的良好候选者 [1]。量化模型后，通常会发生两件事——模型变得更小，运行效率更高。硬件供应商明确允许更快地处理 8 位数据（比 32 位数据快），从而提高吞吐量。更小的模型具有更低的内存占用和功耗 [2]，这对于边缘部署至关重要。
映射函数
映射函数正如您所猜测的——一个将浮点值映射到整数空间的函数。一个常用的映射函数是线性变换，由 给出，其中
是输入，
是 量化参数。
要重新转换回浮点空间，逆函数由 给出。
，它们的差异构成了 量化误差。
量化参数
映射函数由 缩放因子 和 零点
参数化。
简单来说是输入范围与输出范围的比值
其中 [] 是输入的裁剪范围，即允许输入的边界。 [
] 是它映射到的量化输出空间中的范围。对于 8 位量化，输出范围
。
充当偏置，以确保输入空间中的 0 完美映射到量化空间中的 0。
校准
选择输入裁剪范围的过程称为 校准。最简单的技术（也是 PyTorch 中的默认技术）是记录运行中的最小值和最大值，并将它们分配给 和
。TensorRT 还使用熵最小化（KL 散度）、均方误差最小化或输入范围的百分位数。
在 PyTorch 中，Observer 模块（代码）收集输入值的统计信息并计算 qparams 。不同的校准方案会产生不同的量化输出，最好根据经验验证哪种方案最适合您的应用程序和架构（稍后会详细介绍）。
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 时指定仿射或对称方案。请注意，并非所有 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]))
逐张量和逐通道量化方案
量化参数可以为层的整个权重张量作为一个整体计算，也可以为每个通道单独计算。在逐张量中，相同的裁剪范围应用于层中的所有通道。
图 3. 逐通道为每个通道使用一组 qparams。逐张量为整个张量使用相同的 qparams。
对于权重量化，对称逐通道量化提供更好的精度；逐张量量化表现不佳，可能是由于批量归一化折叠导致跨通道的卷积权重方差较大 [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 存储用于量化激活和权重的 Observers 和量化方案。
请务必传递 Observer 类（而不是实例），或可返回 Observer 实例的 callable。使用 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 万个激活。
训练后动态/仅权重（Weight-only）量化
在此方法中，模型的权重是预先量化的；激活在推理期间实时（“动态”）量化。这是所有方法中最简单的一种，它在 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]。下面的示例为了方便起见在校准中使用随机数据——在您的应用程序中使用它将导致不好的 qparams。
图 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)
图 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 产生更高的精度。
(+) 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 个卷积层（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（信号量化噪声比）和均方误差。这种比较分析也可能有助于指导进一步的优化。
图 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. 建议的量化工作流程
注意事项
- 大型（1000 万以上参数）模型对量化误差更具鲁棒性。 [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). 用于高效神经网络推理的量化方法综述。arXiv 预印本 arXiv:2103.13630。
[2] Krishnamoorthi, R. (2018). 量化深度卷积网络以实现高效推理：白皮书。arXiv 预印本 arXiv:1806.08342。
[3] Wu, H., Judd, P., Zhang, X., Isaev, M., & Micikevicius, P. (2020). 深度学习推理的整数量化：原理和实证评估。arXiv 预印本 arXiv:2004.09602。
[4] PyTorch 量化文档