跳转到主要内容
博客

PyTorch 量化简介

在开发机器学习应用时，有效利用服务器端和设备上的计算资源非常重要。为了支持在服务器和边缘设备上更高效的部署，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 提供了适用于大多数用例的默认实现。

我们在 `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` 执行训练后静态量化的三行代码是：`# 为服务器 (x86) 部署设置量化配置 myModel.qconfig = torch.quantization.get_default_config('fbgemm') # 插入观察器 torch.quantization.prepare(myModel, inplace=True) # 校准模型并收集统计信息 # 转换为量化版本 torch.quantization.convert(myModel, inplace=True)`
  3. 量化感知训练量化感知训练 (QAT) 是第三种方法，也是通常能带来最高准确率的方法。在 QAT 中，所有权重和激活在训练的前向和后向传播期间都进行“伪量化”：也就是说，浮点值被四舍五入以模拟 int8 值，但所有计算仍然使用浮点数完成。因此，训练期间的所有权重调整都是在“感知”到模型最终将被量化的事实的情况下进行的；因此，量化后，这种方法通常比其他两种方法产生更高的准确率。
    • PyTorch API:
      • `torch.quantization.prepare_qat` 将伪量化模块插入模型量化中。
      • 与静态量化 API 类似，`torch.quantization.convert` 在训练完成后实际量化模型。
    例如，在端到端示例中，我们加载了一个预训练模型作为 `qat_model`，然后我们简单地使用以下代码执行量化感知训练：`# 为 QAT 指定量化配置 qat_model.qconfig=torch.quantization.get_default_qat_qconfig('fbgemm') # 准备 QAT torch.quantization.prepare_qat(qat_model, inplace=True) # 转换为量化版本，移除 dropout，以检查每个 epoch 的准确性 quantized_model=torch.quantization.convert(qat_model.eval(), inplace=False)`

设备和算子支持

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

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

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 倍，具体取决于硬件平台和基准测试模型。一些示例结果如下：

模型浮点延迟 (ms)量化延迟 (ms)推理性能提升设备注意事项
BERT5813131.8xXeon-D2191 (1.6GHz)批处理大小 = 1，最大序列长度 = 128，单线程，x86-64，动态量化
Resnet-502141032xXeon-D2191 (1.6GHz)单线程，x86-64，静态量化
Mobilenet-v297175.7x三星 S9静态量化，浮点数基于 Caffe2 运行时，未优化

精度结果

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

计算机视觉模型准确性

模型Top-1 准确率（浮点）Top-1 准确率（量化）量化方案
Googlenet69.869.7静态训练后量化
Inception-v377.577.1静态训练后量化
ResNet-1869.869.4静态训练后量化
Resnet-5076.175.9静态训练后量化
ResNext-101 32x8d79.379静态训练后量化
Mobilenet-v271.971.6量化感知训练
Shufflenet-v269.468.4静态训练后量化

语音和自然语言处理模型准确性

模型F1 (GLUEMRPC) 浮点F1 (GLUEMRPC) 量化量化方案
BERT0.9020.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. 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)