量化食谱¶
此食谱演示了如何量化 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 中的子模块名称列表。
警告
动态量化的一个重要限制是,虽然它是最简单的工作流程,如果您没有准备好使用的预训练量化模型,但它目前只支持 nn.Linear 和 nn.LSTM 在 qconfig_spec 中,这意味着您必须使用静态量化或量化感知训练,将在稍后讨论,以量化其他模块,例如 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 格式,然后针对移动应用程序进行优化。有关详细信息,请参阅 移动脚本和优化配方。