PyTorch 上的量化简介

作者：Raghuraman Krishnamoorthi, James Reed, Min Ni, Chris Gottbrath, and Seth Weidman

在开发机器学习应用程序时，高效利用服务器端和设备端计算资源非常重要。为了支持在服务器和边缘设备上更高效地部署，PyTorch 增加了使用熟悉的 eager 模式 Python API 对模型进行量化的支持。

量化利用 8 位整数 (int8) 指令来减小模型大小并加快推理速度（降低延迟），这可能是模型能否达到服务质量目标甚至是否能适应移动设备可用资源的关键。即使资源不是那么受限，它也可能使您能够部署更大、更准确的模型。PyTorch 从 1.3 版本开始支持量化，随着 PyTorch 1.4 版本的发布，我们在 PyTorch torchvision 0.5 库中发布了 ResNet、ResNext、MobileNetV2、GoogleNet、InceptionV3 和 ShuffleNetV2 的量化模型。

这篇博文概述了 PyTorch 上的量化支持及其与 TorchVision 领域库的集成。

什么是量化？

量化是指使用低精度数据（通常是 int8，而非浮点实现）进行计算和内存访问的技术。这可以在几个重要领域实现性能提升

• 模型大小缩小 4 倍；
• 内存带宽减少 2-4 倍；
• 由于节省了内存带宽以及使用 int8 算术计算速度更快，推理速度提升 2-4 倍（具体加速取决于硬件、运行时和模型）。

然而，量化并非没有额外成本。从根本上说，量化意味着引入近似，由此产生的网络的精度会略有下降。这些技术试图最小化全浮点精度与量化精度之间的差距。

我们将量化设计成适应 PyTorch 框架。这意味着

1. PyTorch 具有与量化张量对应的数据类型，这些数据类型共享张量的许多特性。
2. 可以使用量化张量编写内核，就像为浮点张量编写内核一样，以自定义其实现。作为 torch.nn.quantized 和 torch.nn.quantized.dynamic 命名空间的一部分，PyTorch 为常见操作提供量化模块。
3. 量化与其他 PyTorch 功能兼容：量化模型可追踪 (traceable) 且可脚本化 (scriptable)。服务器和移动后端使用的量化方法几乎相同。可以在模型中轻松混合量化和浮点操作。
4. 浮点张量到量化张量的映射可通过用户定义的观察器/伪量化块进行自定义。PyTorch 提供适用于大多数用例的默认实现。

作为 torch.quantization 命名空间中量化工具的一部分，我们在 PyTorch 中开发了三种用于量化神经网络的技术。

PyTorch 1.3 版本开始支持的三种量化模式

1. 动态量化

   PyTorch 支持的最简单的量化方法称为动态量化。这不仅涉及将权重转换为 int8（所有量化变体都会这样做），还涉及在进行计算之前实时将激活值转换为 int8（因此称为“动态”）。计算将因此使用高效的 int8 矩阵乘法和卷积实现进行，从而加快计算速度。但是，激活值是以浮点格式读写到内存的。

   • PyTorch API：我们在 PyTorch 中有一个简单的 API 用于动态量化。torch.quantization.quantize_dynamic 接受一个模型以及一些其他参数，并生成一个量化模型！我们的端到端教程以 BERT 模型为例说明了这一点；虽然该教程很长，包含加载预训练模型和与量化无关的其他概念的部分，但对 BERT 模型进行量化的部分仅仅是

   import torch.quantization
   quantized_model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
   

   • 请查看函数此处的文档以及我们的教程此处和此处的端到端示例。

2. 训练后静态量化

   可以通过将网络转换为使用整数算术和 int8 内存访问来进一步提高性能（延迟）。静态量化会执行额外的步骤，即首先将批量数据馈送通过网络并计算不同激活值的分布（具体来说，这是通过在不同点插入记录这些分布的“观察器”模块来完成的）。此信息用于确定在推理时应如何具体量化不同的激活值（一个简单的技术是将整个激活值范围划分为 256 个级别，但我们也支持更复杂的方法）。重要的是，这个额外的步骤允许我们在操作之间传递量化值，而不是在每次操作之间将这些值转换为浮点数，然后再转回整数，从而显着加速。

   在此版本中，我们支持多项功能，使用户能够优化其静态量化

   1. 观察器：您可以自定义观察器模块，这些模块指定在量化之前如何收集统计信息，以尝试更高级的方法来量化您的数据。
   2. 算子融合：您可以将多个操作融合为一个操作，从而节省内存访问并提高操作的数值精度。
   3. 每通道量化：我们可以对卷积/线性层中的每个输出通道的权重进行独立量化，这可以在几乎相同的速度下带来更高的精度。

   • PyTorch API:

     • 要融合模块，我们有 torch.quantization.fuse_modules
     • 使用 torch.quantization.prepare 插入观察器
     • 最后，量化本身使用 torch.quantization.convert 完成

   我们有一个包含量化端到端示例的教程（这个教程也涵盖了我们的第三种量化方法，量化感知训练），但由于我们简单的 API，对预训练模型 myModel 执行训练后静态量化的三行代码是

   # set quantization config for server (x86)
   deploymentmyModel.qconfig = torch.quantization.get_default_config('fbgemm')
   
   # insert observers
   torch.quantization.prepare(myModel, inplace=True)
   # Calibrate the model and collect statistics
   
   # convert to quantized version
   torch.quantization.convert(myModel, inplace=True)
   

3. 量化感知训练

   量化感知训练 (QAT) 是第三种方法，也是这三种方法中通常能获得最高精度的方法。使用 QAT，在训练的前向和后向传播过程中，所有权重和激活值都经过“伪量化”：也就是说，浮点值被舍入以模仿 int8 值，但所有计算仍然使用浮点数完成。因此，训练期间的所有权重调整都是“感知”到模型最终将被量化的事实而进行的；因此，量化后，这种方法通常比其他两种方法产生更高的精度。

   • PyTorch API:

     • torch.quantization.prepare_qat 将伪量化模块插入到模型量化中。
     • 模仿静态量化 API，一旦训练完成，torch.quantization.convert 实际量化模型。

   例如，在端到端示例中，我们将预训练模型作为 qat_model 加载，然后我们只需使用以下代码执行量化感知训练

   # specify quantization config for QAT
   qat_model.qconfig=torch.quantization.get_default_qat_qconfig('fbgemm')
   
   # prepare QAT
   torch.quantization.prepare_qat(qat_model, inplace=True)
   
   # convert to quantized version, removing dropout, to check for accuracy on each
   epochquantized_model=torch.quantization.convert(qat_model.eval(), inplace=False)
   

设备和算子支持

量化支持仅限于可用算子的子集，具体取决于所使用的方法。有关支持的算子列表，请参阅https://pytorch.ac.cn/docs/stable/quantization.html 上的文档。

可用算子集合和量化数值也取决于用于运行量化模型的后端。目前，量化算子仅在以下后端的 CPU 推理中受支持：x86 和 ARM。量化配置（张量应如何量化）和量化内核（使用量化张量进行的算术）都取决于后端。可以通过以下方式指定后端

import torchbackend='fbgemm'
# 'fbgemm' for server, 'qnnpack' for mobile
my_model.qconfig = torch.quantization.get_default_qconfig(backend)
# prepare and convert model
# Set the backend on which the quantized kernels need to be run
torch.backends.quantized.engine=backend

然而，量化感知训练是在全浮点模式下进行的，并且可以在 GPU 或 CPU 上运行。量化感知训练通常仅在训练后静态或动态量化未能获得足够精度时用于 CNN 模型。这可能发生在为实现小尺寸而高度优化的模型（例如 Mobilenet）上。

在 torchvision 中的集成

我们还为 torchvision 中一些最受欢迎的模型启用了量化：Googlenet、Inception、Resnet、ResNeXt、Mobilenet 和 Shufflenet。我们将这些更改以上述三种形式上游到 torchvision 中

1. 预训练的量化权重，以便您可以立即使用它们。
2. 量化就绪的模型定义，以便您可以进行训练后量化或量化感知训练。
3. 一个用于进行量化感知训练的脚本——虽然该脚本适用于任何这些模型，但正如您将在下面了解到的那样，我们发现只有在 Mobilenet 上实现精度时才需要它。
4. 我们还有一个教程，展示了如何使用其中一个 torchvision 模型进行带量化的迁移学习。

选择方法

选择哪种方案取决于多种因素

1. 模型/目标要求：某些模型可能对量化敏感，需要进行量化感知训练。
2. 算子/后端支持：某些后端需要完全量化的算子。

目前，算子覆盖范围有限，可能会限制下表中列出的选择：下表提供了指导。

性能结果

与浮点实现相比，量化可将模型大小缩小 4 倍，并将速度提升 2 到 3 倍，具体取决于硬件平台和被测试的模型。以下是一些示例结果

精度结果

我们还在 Imagenet 上比较了静态量化模型与浮点模型的精度。对于动态量化，我们在 GLUE 基准测试的 MRPC 上比较了 BERT 的 F1 分数。

计算机视觉模型精度

语音和 NLP 模型精度

总结

要开始在 PyTorch 中量化您的模型，请从PyTorch 网站上的教程开始。如果您正在处理序列数据，请从用于 LSTM 的动态量化或 BERT 的动态量化开始。如果您正在处理图像数据，我们建议您从带量化的迁移学习教程开始。然后您可以探索训练后静态量化。如果您发现训练后量化的精度下降过高，那么尝试量化感知训练。

如果您遇到问题，可以在 discuss.pytorch.org 发帖以获得社区帮助，使用量化分类来解决量化相关问题。

本文由 Raghuraman Krishnamoorthi, James Reed, Min Ni, Chris Gottbrath 和 Seth Weidman 撰写。特别感谢 Jianyu Huang, Lingyi Liu 和 Haixin Liu 提供了本文中包含的量化指标。

延伸阅读:

1. PyTorch 在 Neurips 上的量化演示文稿：(https://research.fb.com/wp-content/uploads/2019/12/2.-Quantization.pptx)
2. 量化张量 (https://github.com/pytorch/pytorch/wiki/ Introducing-Quantized-Tensor)
3. Github 上的量化 RFC (https://github.com/pytorch/pytorch/ issues/18318)