量化秘籍¶
创建时间：2020 年 10 月 26 日 | 最后更新时间：2024 年 3 月 11 日 | 最后验证时间：2024 年 11 月 5 日
本秘籍演示了如何对 PyTorch 模型进行量化，以便在保持与原始模型大致相同精度的情况下，减小模型大小并加快推理速度。量化可以应用于服务器和移动模型的部署，但对于移动端尤为重要甚至关键，因为未经量化的模型大小可能超过 iOS 或 Android 应用允许的限制，导致部署或 OTA 更新耗时过长，并使推理速度过慢，影响用户体验。
引言¶
量化是一种将模型参数中的 32 位浮点数转换为 8 位整数的技术。通过量化，模型大小和内存占用可以减少到原始大小的 1/4，推理速度可以加快约 2-4 倍，同时精度保持大致不变。
对模型进行量化主要有三种方法或工作流程：训练后动态量化、训练后静态量化以及量化感知训练。但是，如果你要使用的模型已经有量化版本，则可以直接使用，无需经过上述三种工作流程中的任何一种。例如，torchvision 库已经包含 MobileNet v2、ResNet 18、ResNet 50、Inception v3、GoogleNet 等模型的量化版本。因此，我们将最后一种方法也算作一种工作流程，尽管它很简单。
注意
量化支持仅适用于有限的一组算子。请参阅此链接了解更多信息。
工作流程¶
使用以下四种工作流程之一对模型进行量化。
1. 使用预训练的量化 MobileNet v2¶
要获取 MobileNet v2 量化模型，只需执行
import torchvision
model_quantized = torchvision.models.quantization.mobilenet_v2(pretrained=True, quantize=True)
要比较非量化 MobileNet v2 模型与其量化版本的大小差异
model = torchvision.models.mobilenet_v2(pretrained=True)
import os
import torch
def print_model_size(mdl):
torch.save(mdl.state_dict(), "tmp.pt")
print("%.2f MB" %(os.path.getsize("tmp.pt")/1e6))
os.remove('tmp.pt')
print_model_size(model)
print_model_size(model_quantized)
输出将是
14.27 MB
3.63 MB
2. 训练后动态量化¶
要应用动态量化，它将模型中的所有权重从 32 位浮点数转换为 8 位整数，但在对激活执行计算之前才将激活转换为 int8，只需调用 torch.quantization.quantize_dynamic
model_dynamic_quantized = torch.quantization.quantize_dynamic(
model, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8
)
其中 qconfig_spec 指定了要对模型中应用量化的子模块名称列表。
警告
动态量化一个重要的限制是，虽然在你没有现成的预训练量化模型时它是最简单的工作流程，但它目前在 qconfig_spec 中仅支持 nn.Linear 和 nn.LSTM，这意味着你需要使用稍后讨论的静态量化或量化感知训练来量化其他模块，例如 nn.Conv2d。
quantize_dynamic API 调用的完整文档在此处：此处。使用训练后动态量化的其他三个示例是Bert 示例、一个 LSTM 模型示例以及另一个LSTM 演示示例。
3. 训练后静态量化¶
此方法预先将权重和激活都转换为 8 位整数，因此在推理过程中不会像动态量化那样对激活进行即时转换。虽然训练后静态量化可以显著提高推理速度并减小模型大小，但与训练后动态量化相比，此方法可能会更大地降低原始模型的精度。
要对模型应用静态量化，运行以下代码
backend = "qnnpack"
model.qconfig = torch.quantization.get_default_qconfig(backend)
torch.backends.quantized.engine = backend
model_static_quantized = torch.quantization.prepare(model, inplace=False)
model_static_quantized = torch.quantization.convert(model_static_quantized, inplace=False)
在此之后，运行 print_model_size(model_static_quantized) 显示静态量化模型大小为 3.98MB。
完整的模型定义和静态量化示例在此处：此处。专门的静态量化教程在此处：此处。
注意
要使模型在通常采用 arm 架构的移动设备上运行，backend 需要使用 qnnpack；要在采用 x86 架构的计算机上运行模型，使用 x86`（旧的 fbgemm 仍然可用，但推荐默认使用 `x86`）。
4. 量化感知训练¶
量化感知训练在模型训练过程中对所有权重和激活插入伪量化，从而获得比训练后量化方法更高的推理精度。它通常用于 CNN 模型。
要启用模型进行量化感知训练，请在模型定义的 __init__ 方法中定义一个 QuantStub 和一个 DeQuantStub，用于将张量从浮点类型转换为量化类型，反之亦然
self.quant = torch.quantization.QuantStub()
self.dequant = torch.quantization.DeQuantStub()
然后，在模型定义的 forward 方法的开头和结尾调用 x = self.quant(x) 和 x = self.dequant(x)。
要进行量化感知训练，使用以下代码片段
model.qconfig = torch.quantization.get_default_qat_qconfig(backend)
model_qat = torch.quantization.prepare_qat(model, inplace=False)
# quantization aware training goes here
model_qat = torch.quantization.convert(model_qat.eval(), inplace=False)
预训练的量化模型也可以用于量化感知的迁移学习，使用上面显示的相同的 quant 和 dequant 调用。请参阅此处获取完整示例。
使用上述步骤之一生成量化模型后，在模型可用于移动设备之前，需要进一步将其转换为 TorchScript 格式，然后针对移动应用进行优化。有关详细信息，请参阅脚本和优化用于移动设备的秘籍。
