快捷方式

量化食谱

此食谱演示了如何量化 PyTorch 模型,以便它能够以更小的尺寸和更快的推理速度运行,同时保持与原始模型相同的准确率。量化可以应用于服务器和移动模型部署,但在移动设备上尤为重要甚至至关重要,因为未量化的模型的尺寸可能会超出 iOS 或 Android 应用程序允许的限制,导致部署或 OTA 更新耗时过长,并且推理速度太慢,无法提供良好的用户体验。

简介

量化是一种将模型参数中的 32 位浮点数转换为 8 位整数的技术。使用量化,模型大小和内存占用可以减少到原始大小的 1/4,并且推理速度可以提高约 2-4 倍,同时保持相同的准确率。

总共有三种量化模型的方法或工作流程:训练后动态量化、训练后静态量化和量化感知训练。但如果要使用的模型已经存在量化版本,则可以使用它,无需执行上述任何三种工作流程。例如,torchvision 库已经包含了 MobileNet v2、ResNet 18、ResNet 50、Inception v3、GoogleNet 等模型的量化版本。因此,我们将最后一种方法作为另一种工作流程,尽管它很简单。

注意

量化支持仅适用于有限的运算符集。有关详细信息,请参阅 此处

先决条件

PyTorch 1.6.0 或 1.7.0

torchvision 0.6.0 或 0.7.0

工作流程

使用以下四种工作流程之一来量化模型。

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.Linearnn.LSTMqconfig_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)

有关量化感知训练的更详细示例,请参阅 此处此处

预训练的量化模型也可以用于量化感知迁移学习,使用上面显示的相同的 quantdequant 调用。请参阅 此处 以获取完整的示例。

使用上述步骤之一生成量化模型后,在模型可以用于在移动设备上运行之前,需要将其进一步转换为 TorchScript 格式,然后针对移动应用程序进行优化。有关详细信息,请参阅 移动脚本和优化配方

了解更多信息

有关量化不同工作流程的更多信息,请参阅 此处此处

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源