量化配方¶
创建于：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 指定了 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架构的移动设备上运行，您需要使用qnnpack作为backend；要在具有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格式，然后针对移动应用程序进行优化。详情请参阅Script and Optimize for Mobile recipe。
