PyTorch 量化入门

由 Raghuraman Krishnamoorthi、James Reed、Min Ni、Chris Gottbrath 和 Seth Weidman 撰写

在开发机器学习应用时，高效利用服务器端和设备端计算资源非常重要。为了支持在服务器和边缘设备上更高效的部署，PyTorch 添加了对使用熟悉的即时模式 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. 可以使用量化张量编写内核，就像为浮点张量编写的内核一样，以自定义其实现。PyTorch 支持量化模块，用于常见的操作，作为 torch.nn.quantized 和 torch.nn.quantized.dynamic 命名空间的一部分。
3. 量化与 PyTorch 的其他部分兼容：量化模型可追踪和可脚本化。量化方法对于服务器和移动后端几乎相同。可以在模型中轻松混合量化和浮点运算。
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://maskerprc.github.io/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）