量化配方¶
创建于：2020 年 10 月 26 日 | 最后更新：2024 年 3 月 11 日 | 最后验证：2024 年 11 月 05 日
本配方演示了如何量化 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 指定要对其应用量化的 model 中的子模块名称列表。
警告
动态量化的一个重要限制是，如果您没有可供使用的预训练量化模型，那么虽然它是最简单的工作流程，但它目前仅支持 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 格式，然后针对移动应用程序进行优化。有关详细信息，请参阅脚本并优化移动设备配方。
