量化是一种廉价且简单的方法，可以使您的 DNN 运行更快，并降低内存要求。PyTorch 提供了几种不同的模型量化方法。在这篇博文中，我们将（快速）介绍深度学习中量化的基础知识，然后实际看看每种技术是什么样的。最后，我们将以文献中关于在您的工作流程中使用量化的建议结束。
图 1. PyTorch ❤️ 量化
目录
量化基础
如果有人问你时间，你不会回答“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))
后端引擎
目前，量化运算符在 x86 机器上通过 FBGEMM 后端运行，或在 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（代码）命名元组存储用于量化激活和权重的观察器和量化方案。
请务必传递 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 模式 vs FX 图模式）
- 如果用于量化激活（层输出）的 qparams 是为所有输入预先计算的，还是每次输入时重新计算的（静态 vs 动态），
- 如果 qparams 是在重新训练时计算的，还是不重新训练时计算的（量化感知训练 vs 训练后量化）
FX 图模式会自动融合符合条件的模块，插入 Quant/DeQuant 存根，校准模型并返回量化模块——所有这些都只需两次方法调用——但仅适用于可 符号跟踪 的网络。下面的示例包含使用 Eager 模式和 FX 图模式的调用以进行比较。
在 DNN 中，符合量化条件的候选者是 FP32 权重（层参数）和激活（层输出）。量化权重可减小模型大小。量化激活通常会加速推理。
例如，50 层 ResNet 网络具有约 2600 万个权重参数，并在前向传播中计算约 1600 万个激活。
训练后动态/仅权重（Post-Training Dynamic/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，我们可以轻松创建自定义 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). 用于高效神经网络推理的量化方法综述。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 量化文档