作者：Raghuraman Krishnamoorthi、James Reed、Min Ni、Chris Gottbrath 和 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. 可以像编写浮点张量的内核一样，使用量化张量编写内核，以自定义其实现。PyTorch 在 torch.nn.quantized 和 torch.nn.quantized.dynamic 命名空间中支持常用操作的量化模块。
  3. 量化与 PyTorch 的其余部分兼容：量化模型是可追踪和可脚本化的。服务器和移动后端采用的量化方法几乎相同。可以在模型中轻松混合量化和浮点运算。
  4. 浮点张量到量化张量的映射可以通过用户定义的观察者/伪量化块进行自定义。PyTorch 提供了默认实现，这些实现应该适用于大多数用例。

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

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. 运算符/后端支持：某些后端需要完全量化的运算符。

目前，运算符覆盖范围有限，可能会限制下表列出的选择：下表提供了一个指南。

模型类型 首选方案 原因
LSTM/RNN 动态量化 吞吐量主要受权重计算/内存带宽限制
BERT/Transformer 动态量化 吞吐量主要受权重计算/内存带宽限制
CNN 静态量化 吞吐量受激活内存带宽限制
CNN 量化感知训练 在静态量化无法实现精度的情况下

性能结果

量化使模型大小减少 4 倍，并且与浮点实现相比，速度提高了 2 倍到 3 倍，具体取决于硬件平台和要进行基准测试的模型。一些示例结果如下

模型 浮点延迟（毫秒） 量化延迟（毫秒） 推理性能提升 设备 注释
BERT 581 313 1.8 倍 Xeon-D2191 (1.6GHz) 批大小 = 1，最大序列长度 = 128，单线程，x86-64，动态量化
Resnet-50 214 103 2 倍 Xeon-D2191 (1.6GHz) 单线程，x86-64，静态量化
Mobilenet-v2 97 17 5.7 倍 三星 S9 静态量化，浮点数基于 Caffe2 运行时且未优化

精度结果

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

计算机视觉模型精度

模型 Top-1 精度（浮点） Top-1 精度（量化） 量化方案
Googlenet 69.8 69.7 静态训练后量化
Inception-v3 77.5 77.1 静态训练后量化
ResNet-18 69.8 69.4 静态训练后量化
Resnet-50 76.1 75.9 静态训练后量化
ResNext-101 32x8d 79.3 79 静态训练后量化
Mobilenet-v2 71.9 71.6 量化感知训练
Shufflenet-v2 69.4 68.4 静态训练后量化

语音和 NLP 模型精度

模型 F1 (GLUEMRPC) 浮点 F1 (GLUEMRPC) 量化 量化方案
BERT 0.902 0.895 动态量化

结论

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

如果您遇到问题，可以在 discuss.pytorch.org 上发帖寻求社区帮助，量化相关问题请使用量化类别。

这篇文章由 Raghuraman Krishnamoorthi、James Reed、Min Ni、Chris Gottbrath 和 Seth Weidman 撰写。特别感谢 Jianyu Huang、Lingyi Liu 和 Haixin Liu 制作了本文中包含的量化指标。

延伸阅读:

  1. Neurips 上的 PyTorch 量化演示：(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)