跳转到主要内容
博客

PyTorch 量化简介

在开发机器学习应用程序时，有效利用服务器端和设备上的计算资源非常重要。为了支持在服务器和边缘设备上进行更高效的部署，PyTorch 添加了对模型量化的支持，使用的是熟悉的 Eager Mode 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` 上执行训练后静态量化的三行代码是：`# 为服务器 (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的文档。

可用操作符集和量化数值也取决于用于运行量化模型的后端。目前，量化操作符仅支持以下后端中的 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 倍。一些示例结果如下：

模型浮点延迟 (ms)量化延迟 (ms)推理性能增益设备注意事项
BERT5813131.8 倍Xeon-D2191 (1.6GHz)批大小 = 1，最大序列长度 = 128，单线程，x86-64，动态量化
Resnet-502141032 倍Xeon-D2191 (1.6GHz)单线程，x86-64，静态量化
Mobilenet-v297175.7 倍三星 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训练后静态量化

语音和 NLP 模型精度

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