量化

警告

量化处于测试阶段，可能会发生变化。

量化简介

量化指的是使用比浮点精度更低的位宽执行计算和存储张量的技术。量化模型在具有降低精度的张量上执行某些或所有操作，而不是使用全精度（浮点）值。这使得模型表示更加紧凑，并且可以在许多硬件平台上使用高性能矢量化操作。PyTorch 支持 INT8 量化，而典型的 FP32 模型则允许模型大小减少 4 倍，内存带宽需求减少 4 倍。与 FP32 计算相比，INT8 计算的硬件支持通常快 2 到 4 倍。量化主要是一种加速推理的技术，并且仅支持量化算子的前向传播。

PyTorch 支持多种量化深度学习模型的方法。在大多数情况下，模型以 FP32 训练，然后转换为 INT8。此外，PyTorch 还支持量化感知训练，该训练使用伪量化模块在正向和反向传播中模拟量化误差。请注意，整个计算是在浮点中执行的。在量化感知训练结束时，PyTorch 提供转换函数以将训练后的模型转换为较低的精度。

在较低级别，PyTorch 提供了一种表示量化张量并对其执行操作的方法。它们可用于直接构建执行全部或部分低精度计算的模型。提供了更高级别的 API，这些 API 集成了将 FP32 模型转换为较低精度且精度损失最小的典型工作流程。

量化 API 摘要

PyTorch 提供三种不同的量化模式：急切模式量化、FX 图模式量化（维护）和 PyTorch 2 导出量化。

急切模式量化是一个测试功能。用户需要手动执行融合并指定量化和反量化发生的位置，并且它仅支持模块，而不支持函数。

FX 图模式量化是 PyTorch 中的自动化量化工作流程，目前它是一个原型功能，由于我们有 PyTorch 2 导出量化，因此它处于维护模式。它通过添加对函数的支持和自动化量化过程来改进急切模式量化，尽管人们可能需要重构模型以使模型与 FX 图模式量化兼容（可以使用 torch.fx 进行符号跟踪）。请注意，FX 图模式量化预计不会在任意模型上运行，因为模型可能无法进行符号跟踪，我们将其集成到域库（如 torchvision）中，用户将能够使用 FX 图模式量化来量化类似于受支持域库中的模型。对于任意模型，我们将提供一般准则，但要使其真正起作用，用户可能需要熟悉 torch.fx，尤其是在如何使模型能够进行符号跟踪方面。

PyTorch 2 导出量化是新的完整图模式量化工作流程，作为 PyTorch 2.1 中的原型功能发布。使用 PyTorch 2，我们正在转向一个更好的完整程序捕获解决方案（torch.export），因为它与 torch.fx.symbolic_trace（在 14K 个模型上为 72.7%）相比，可以捕获更高百分比（在 14K 个模型上为 88.8%）的模型，FX 图模式量化使用的程序捕获解决方案。torch.export 在某些 Python 结构方面仍然存在局限性，并且需要用户参与来支持导出模型中的动态性，但总的来说，它是对先前程序捕获解决方案的改进。PyTorch 2 导出量化是为 torch.export 捕获的模型构建的，同时兼顾建模用户和后端开发人员的灵活性和生产力。主要功能包括：(1). 可编程 API 用于配置模型的量化方式，可以扩展到更多用例 (2). 简化的建模用户和后端开发人员的 UX，因为他们只需要与单个对象（量化器）交互来表达用户关于如何量化模型以及后端支持的意图。(3). 可选的参考量化模型表示，可以使用与硬件中发生的实际量化计算更接近的整数运算来表示量化计算。

鼓励量化的新用户首先尝试 PyTorch 2 导出量化，如果效果不佳，用户可以尝试急切模式量化。

下表比较了急切模式量化、FX 图模式量化和 PyTorch 2 导出量化之间的差异

	急切模式量化	FX 图模式量化	PyTorch 2 导出量化
发布状态	测试版	原型（维护）	原型
算子融合	手动	自动	自动
量化/反量化放置	手动	自动	自动
量化模块	支持	支持	支持
量化函数/Torch 操作	手动	自动	支持
对自定义的支持	有限支持	完全支持	完全支持
量化模式支持	训练后量化：静态、动态、仅权重 量化感知训练：静态	训练后量化：静态、动态、仅权重 量化感知训练：静态	由后端特定量化器定义
输入/输出模型类型	torch.nn.Module	torch.nn.Module（可能需要一些重构才能使模型与 FX 图模式量化兼容）	torch.fx.GraphModule（由 torch.export 捕获）

支持三种类型的量化

1. 动态量化（权重量化，激活值以浮点数读取/存储，并为计算量化）

2. 静态量化（权重量化，激活值量化，训练后需要校准）

3. 静态量化感知训练（权重量化，激活量化，训练期间对量化数值建模）

请参阅我们的 PyTorch 量化简介 博客文章，以更全面地了解这些量化类型之间的权衡。

算子的覆盖范围在动态量化和静态量化之间有所不同，并在下表中进行了说明。

	静态量化	动态量化
nn.Linear nn.Conv1d/2d/3d	Y Y	Y N
nn.LSTM nn.GRU	Y（通过 自定义模块） N	Y Y
nn.RNNCell nn.GRUCell nn.LSTMCell	N N N	Y Y Y
nn.EmbeddingBag	Y（激活处于 fp32）	Y
nn.Embedding	Y	Y
nn.MultiheadAttention	Y（通过自定义模块）	不支持
激活	广泛支持	保持不变，计算保持在 fp32

急切模式量化

有关量化流程的总体介绍，包括不同类型的量化，请查看 通用量化流程。

训练后动态量化

这是最简单的量化应用形式，其中权重在预先量化，但激活在推理期间动态量化。这用于模型执行时间主要由从内存加载权重而不是计算矩阵乘法主导的情况。对于具有小批量大小的 LSTM 和 Transformer 类型模型，这是正确的。

图表

# original model
# all tensors and computations are in floating point
previous_layer_fp32 -- linear_fp32 -- activation_fp32 -- next_layer_fp32
                 /
linear_weight_fp32

# dynamically quantized model
# linear and LSTM weights are in int8
previous_layer_fp32 -- linear_int8_w_fp32_inp -- activation_fp32 -- next_layer_fp32
                     /
   linear_weight_int8

PTDQ API 示例

import torch

# define a floating point model
class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(4, 4)

    def forward(self, x):
        x = self.fc(x)
        return x

# create a model instance
model_fp32 = M()
# create a quantized model instance
model_int8 = torch.ao.quantization.quantize_dynamic(
    model_fp32,  # the original model
    {torch.nn.Linear},  # a set of layers to dynamically quantize
    dtype=torch.qint8)  # the target dtype for quantized weights

# run the model
input_fp32 = torch.randn(4, 4, 4, 4)
res = model_int8(input_fp32)

要了解有关动态量化的更多信息，请参阅我们的 动态量化教程。

训练后静态量化

训练后静态量化 (PTQ 静态) 量化模型的权重和激活。它尽可能地将激活融合到前面的层中。它需要使用代表性数据集进行校准，以确定激活的最佳量化参数。当内存带宽和计算节省都很重要时，通常使用训练后静态量化，CNN 是一个典型的用例。

在应用训练后静态量化之前，我们可能需要修改模型。请参阅 急切模式静态量化的模型准备。

图表

# original model
# all tensors and computations are in floating point
previous_layer_fp32 -- linear_fp32 -- activation_fp32 -- next_layer_fp32
                    /
    linear_weight_fp32

# statically quantized model
# weights and activations are in int8
previous_layer_int8 -- linear_with_activation_int8 -- next_layer_int8
                    /
  linear_weight_int8

PTSQ API 示例

import torch

# define a floating point model where some layers could be statically quantized
class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # QuantStub converts tensors from floating point to quantized
        self.quant = torch.ao.quantization.QuantStub()
        self.conv = torch.nn.Conv2d(1, 1, 1)
        self.relu = torch.nn.ReLU()
        # DeQuantStub converts tensors from quantized to floating point
        self.dequant = torch.ao.quantization.DeQuantStub()

    def forward(self, x):
        # manually specify where tensors will be converted from floating
        # point to quantized in the quantized model
        x = self.quant(x)
        x = self.conv(x)
        x = self.relu(x)
        # manually specify where tensors will be converted from quantized
        # to floating point in the quantized model
        x = self.dequant(x)
        return x

# create a model instance
model_fp32 = M()

# model must be set to eval mode for static quantization logic to work
model_fp32.eval()

# attach a global qconfig, which contains information about what kind
# of observers to attach. Use 'x86' for server inference and 'qnnpack'
# for mobile inference. Other quantization configurations such as selecting
# symmetric or asymmetric quantization and MinMax or L2Norm calibration techniques
# can be specified here.
# Note: the old 'fbgemm' is still available but 'x86' is the recommended default
# for server inference.
# model_fp32.qconfig = torch.ao.quantization.get_default_qconfig('fbgemm')
model_fp32.qconfig = torch.ao.quantization.get_default_qconfig('x86')

# Fuse the activations to preceding layers, where applicable.
# This needs to be done manually depending on the model architecture.
# Common fusions include `conv + relu` and `conv + batchnorm + relu`
model_fp32_fused = torch.ao.quantization.fuse_modules(model_fp32, [['conv', 'relu']])

# Prepare the model for static quantization. This inserts observers in
# the model that will observe activation tensors during calibration.
model_fp32_prepared = torch.ao.quantization.prepare(model_fp32_fused)

# calibrate the prepared model to determine quantization parameters for activations
# in a real world setting, the calibration would be done with a representative dataset
input_fp32 = torch.randn(4, 1, 4, 4)
model_fp32_prepared(input_fp32)

# Convert the observed model to a quantized model. This does several things:
# quantizes the weights, computes and stores the scale and bias value to be
# used with each activation tensor, and replaces key operators with quantized
# implementations.
model_int8 = torch.ao.quantization.convert(model_fp32_prepared)

# run the model, relevant calculations will happen in int8
res = model_int8(input_fp32)

要了解有关静态量化的更多信息，请参阅 静态量化教程。

静态量化的量化感知训练

量化感知训练 (QAT) 在训练期间模拟量化的影响，与其他量化方法相比，可以获得更高的精度。我们可以对静态、动态或仅权重量化进行 QAT。在训练期间，所有计算都在浮点中完成，fake_quant 模块通过钳位和舍入模拟量化的影响，以模拟 INT8 的影响。模型转换后，权重和激活被量化，并且激活尽可能地融合到前面的层中。它通常与 CNN 一起使用，并且与静态量化相比具有更高的精度。

在应用训练后静态量化之前，我们可能需要修改模型。请参阅 急切模式静态量化的模型准备。

图表

# original model
# all tensors and computations are in floating point
previous_layer_fp32 -- linear_fp32 -- activation_fp32 -- next_layer_fp32
                      /
    linear_weight_fp32

# model with fake_quants for modeling quantization numerics during training
previous_layer_fp32 -- fq -- linear_fp32 -- activation_fp32 -- fq -- next_layer_fp32
                           /
   linear_weight_fp32 -- fq

# quantized model
# weights and activations are in int8
previous_layer_int8 -- linear_with_activation_int8 -- next_layer_int8
                     /
   linear_weight_int8

QAT API 示例

import torch

# define a floating point model where some layers could benefit from QAT
class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # QuantStub converts tensors from floating point to quantized
        self.quant = torch.ao.quantization.QuantStub()
        self.conv = torch.nn.Conv2d(1, 1, 1)
        self.bn = torch.nn.BatchNorm2d(1)
        self.relu = torch.nn.ReLU()
        # DeQuantStub converts tensors from quantized to floating point
        self.dequant = torch.ao.quantization.DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.dequant(x)
        return x

# create a model instance
model_fp32 = M()

# model must be set to eval for fusion to work
model_fp32.eval()

# attach a global qconfig, which contains information about what kind
# of observers to attach. Use 'x86' for server inference and 'qnnpack'
# for mobile inference. Other quantization configurations such as selecting
# symmetric or asymmetric quantization and MinMax or L2Norm calibration techniques
# can be specified here.
# Note: the old 'fbgemm' is still available but 'x86' is the recommended default
# for server inference.
# model_fp32.qconfig = torch.ao.quantization.get_default_qconfig('fbgemm')
model_fp32.qconfig = torch.ao.quantization.get_default_qat_qconfig('x86')

# fuse the activations to preceding layers, where applicable
# this needs to be done manually depending on the model architecture
model_fp32_fused = torch.ao.quantization.fuse_modules(model_fp32,
    [['conv', 'bn', 'relu']])

# Prepare the model for QAT. This inserts observers and fake_quants in
# the model needs to be set to train for QAT logic to work
# the model that will observe weight and activation tensors during calibration.
model_fp32_prepared = torch.ao.quantization.prepare_qat(model_fp32_fused.train())

# run the training loop (not shown)
training_loop(model_fp32_prepared)

# Convert the observed model to a quantized model. This does several things:
# quantizes the weights, computes and stores the scale and bias value to be
# used with each activation tensor, fuses modules where appropriate,
# and replaces key operators with quantized implementations.
model_fp32_prepared.eval()
model_int8 = torch.ao.quantization.convert(model_fp32_prepared)

# run the model, relevant calculations will happen in int8
res = model_int8(input_fp32)

要了解有关量化感知训练的更多信息，请参阅 QAT 教程。

急切模式静态量化的模型准备

目前有必要在急切模式量化之前对模型定义进行一些修改。这是因为当前量化是基于模块的。具体来说，对于所有量化技术，用户都需要

1. 将任何需要输出重新量化（因此具有其他参数）的操作从函数式转换为模块形式（例如，使用 torch.nn.ReLU 而不是 torch.nn.functional.relu）。

2. 指定模型的哪些部分需要量化，方法是在子模块上分配 .qconfig 属性或指定 qconfig_mapping。例如，设置 model.conv1.qconfig = None 表示 model.conv 层不会被量化，并且设置 model.linear1.qconfig = custom_qconfig 表示 model.linear1 的量化设置将使用 custom_qconfig 而不是全局 qconfig。

对于量化激活的静态量化技术，用户还需要执行以下操作

1. 指定激活量化和反量化的位置。这是使用 QuantStub 和 DeQuantStub 模块完成的。

2. 使用 FloatFunctional 将需要针对量化进行特殊处理的张量操作包装到模块中。例如，操作如 add 和 cat 需要特殊处理以确定输出量化参数。

3. 融合模块：将操作/模块合并到单个模块中以获得更高的精度和性能。这是使用 fuse_modules() API 完成的，该 API 接收要融合的模块列表。我们目前支持以下融合：[Conv，Relu]，[Conv，BatchNorm]，[Conv，BatchNorm，Relu]，[Linear，Relu]

（原型 - 维护模式）FX 图模式量化

训练后量化中有多种量化类型（仅权重、动态和静态），配置是通过 qconfig_mapping（prepare_fx 函数的参数）完成的。

FXPTQ API 示例

import torch
from torch.ao.quantization import (
  get_default_qconfig_mapping,
  get_default_qat_qconfig_mapping,
  QConfigMapping,
)
import torch.ao.quantization.quantize_fx as quantize_fx
import copy

model_fp = UserModel()

#
# post training dynamic/weight_only quantization
#

# we need to deepcopy if we still want to keep model_fp unchanged after quantization since quantization apis change the input model
model_to_quantize = copy.deepcopy(model_fp)
model_to_quantize.eval()
qconfig_mapping = QConfigMapping().set_global(torch.ao.quantization.default_dynamic_qconfig)
# a tuple of one or more example inputs are needed to trace the model
example_inputs = (input_fp32)
# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, example_inputs)
# no calibration needed when we only have dynamic/weight_only quantization
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# post training static quantization
#

model_to_quantize = copy.deepcopy(model_fp)
qconfig_mapping = get_default_qconfig_mapping("qnnpack")
model_to_quantize.eval()
# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, example_inputs)
# calibrate (not shown)
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# quantization aware training for static quantization
#

model_to_quantize = copy.deepcopy(model_fp)
qconfig_mapping = get_default_qat_qconfig_mapping("qnnpack")
model_to_quantize.train()
# prepare
model_prepared = quantize_fx.prepare_qat_fx(model_to_quantize, qconfig_mapping, example_inputs)
# training loop (not shown)
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# fusion
#
model_to_quantize = copy.deepcopy(model_fp)
model_fused = quantize_fx.fuse_fx(model_to_quantize)

请遵循以下教程以了解有关 FX 图模式量化的更多信息

• 使用 FX 图模式量化的用户指南

• FX 图模式训练后静态量化

• FX 图模式训练后动态量化

（原型）PyTorch 2 导出量化

API 示例

import torch
from torch.ao.quantization.quantize_pt2e import prepare_pt2e
from torch._export import capture_pre_autograd_graph
from torch.ao.quantization.quantizer import (
    XNNPACKQuantizer,
    get_symmetric_quantization_config,
)

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(5, 10)

   def forward(self, x):
       return self.linear(x)

# initialize a floating point model
float_model = M().eval()

# define calibration function
def calibrate(model, data_loader):
    model.eval()
    with torch.no_grad():
        for image, target in data_loader:
            model(image)

# Step 1. program capture
# NOTE: this API will be updated to torch.export API in the future, but the captured
# result should mostly stay the same
m = capture_pre_autograd_graph(m, *example_inputs)
# we get a model with aten ops

# Step 2. quantization
# backend developer will write their own Quantizer and expose methods to allow
# users to express how they
# want the model to be quantized
quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config())
# or prepare_qat_pt2e for Quantization Aware Training
m = prepare_pt2e(m, quantizer)

# run calibration
# calibrate(m, sample_inference_data)
m = convert_pt2e(m)

# Step 3. lowering
# lower to target backend

请遵循以下教程开始使用 PyTorch 2 导出量化

建模用户

• PyTorch 2 导出训练后量化

• 通过 Inductor 使用 X86 后端进行 PyTorch 2 导出训练后量化

• PyTorch 2 导出量化感知训练

后端开发人员（也请查看所有建模用户文档）

• 如何为 PyTorch 2 导出量化编写量化器

量化堆栈

量化是将浮点模型转换为量化模型的过程。因此，在高级别上，量化堆栈可以分为两部分：1）量化模型的构建块或抽象 2）将浮点模型转换为量化模型的量化流程的构建块或抽象

量化模型

量化张量

为了在 PyTorch 中进行量化，我们需要能够在张量中表示量化数据。量化张量允许存储量化数据（表示为 int8/uint8/int32）以及量化参数，如比例和零点。量化张量允许进行许多有用的操作，使量化算术变得容易，此外还允许以量化格式序列化数据。

PyTorch 支持每张量和每通道的对称和非对称量化。每张量意味着张量中的所有值都以相同的方式使用相同的量化参数进行量化。每通道意味着对于每个维度，通常是张量的通道维度，张量中的值使用不同的量化参数进行量化。这使得将张量转换为量化值的误差更小，因为异常值只会影响它所在的通道，而不是整个张量。

映射是通过使用以下方法转换浮点张量来执行的

请注意，我们确保浮点中的零在量化后没有误差，从而确保诸如填充之类的操作不会导致额外的量化误差。

以下是量化张量的一些关键属性

• QScheme（torch.qscheme）：一个枚举，指定我们量化张量的方式

  • torch.per_tensor_affine

  • torch.per_tensor_symmetric

  • torch.per_channel_affine

  • torch.per_channel_symmetric

• dtype（torch.dtype）：量化张量的数据类型

  • torch.quint8

  • torch.qint8

  • torch.qint32

  • torch.float16

• 量化参数（根据 QScheme 变化）：所选量化方式的参数

  • torch.per_tensor_affine 将具有以下量化参数

    • scale（浮点数）

    • zero_point（整数）

  • torch.per_channel_affine 将具有以下量化参数

    • per_channel_scales（浮点数列表）

    • per_channel_zero_points（整数列表）

    • axis（整数）

量化和反量化

模型的输入和输出是浮点张量，但量化模型中的激活是量化的，因此我们需要算子在浮点和量化张量之间进行转换。

• 量化（浮点 -> 量化）

  • torch.quantize_per_tensor(x, scale, zero_point, dtype)

  • torch.quantize_per_channel(x, scales, zero_points, axis, dtype)

  • torch.quantize_per_tensor_dynamic(x, dtype, reduce_range)

  • to(torch.float16)

• 反量化（量化 -> 浮点）

  • quantized_tensor.dequantize() - 在 torch.float16 张量上调用 dequantize 将把张量转换回 torch.float

  • torch.dequantize(x)

量化算子/模块

• 量化算子是将量化张量作为输入并输出量化张量的算子。

• 量化模块是执行量化操作的 PyTorch 模块。它们通常定义用于线性、卷积等加权操作。

量化引擎

当执行量化模型时，qengine（torch.backends.quantized.engine）指定要用于执行的后端。确保 qengine 在量化激活和权重的值范围方面与量化模型兼容非常重要。

量化流程

观察器和 FakeQuantize

• 观察器是用于

  • 收集张量统计信息，例如通过观察器传递的张量的最小值和最大值

  • 并根据收集到的张量统计信息计算量化参数

• FakeQuantize 是用于

  • 模拟网络中张量的量化（执行量化/反量化）

  • 它可以根据从观察者收集的统计数据计算量化参数，或者也可以学习量化参数。

QConfig

• QConfig 是一个命名元组，包含 Observer 或 FakeQuantize 模块类，可以通过 qscheme、dtype 等进行配置。它用于配置如何观察某个算子。

  • 算子/模块的量化配置

    • 不同类型的 Observer/FakeQuantize

    • 数据类型

    • 量化方案

    • quant_min/quant_max：可用于模拟低精度张量。

  • 目前支持激活和权重的配置。

  • 我们根据为给定算子或模块配置的 qconfig 插入输入/权重/输出观察器。

通用量化流程

一般来说，流程如下：

• 准备

  • 根据用户指定的 qconfig 插入 Observer/FakeQuantize 模块。

• 校准/训练（取决于训练后量化或量化感知训练）。

  • 允许观察器收集统计数据或 FakeQuantize 模块学习量化参数。

• 转换

  • 将校准/训练后的模型转换为量化模型。

量化有不同的模式，可以从两个方面进行分类：

从应用量化流程的位置来看，我们有：

1. 训练后量化（在训练后应用量化，量化参数基于样本校准数据计算）。

2. 量化感知训练（在训练期间模拟量化，以便可以使用训练数据将量化参数与模型一起学习）。

从如何量化算子的角度来看，我们可以有：

• 仅权重量化（仅权重被静态量化）。

• 动态量化（权重被静态量化，激活被动态量化）。

• 静态量化（权重和激活都被静态量化）。

我们可以在同一个量化流程中混合不同的算子量化方式。例如，我们可以进行训练后量化，其中包含静态和动态量化的算子。

量化支持矩阵

量化模式支持

	量化模式		数据集需求	最适合	准确率	备注
训练后量化	动态/仅权重量化	激活动态量化（fp16、int8）或未量化，权重静态量化（fp16、int8、int4）	无	LSTM、MLP、Embedding、Transformer	良好	易于使用，当性能受权重计算或内存限制时，接近静态量化。
	静态量化	激活和权重静态量化（int8）	校准数据集	CNN	良好	提供最佳性能，可能会对准确性产生重大影响，适用于仅支持 int8 计算的硬件。
量化感知训练	动态量化	激活和权重被伪量化。	微调数据集	MLP、Embedding	最佳	目前支持有限。
	静态量化	激活和权重被伪量化。	微调数据集	CNN、MLP、Embedding	最佳	通常在静态量化导致精度下降时使用，用于缩小精度差距。

请参阅我们的Pytorch 量化简介 博客文章，以更全面地了解这些量化类型之间的权衡。

量化流程支持

PyTorch 提供两种量化模式：急切模式量化和 FX 图模式量化。

急切模式量化是一个测试功能。用户需要手动执行融合并指定量化和反量化发生的位置，并且它仅支持模块，而不支持函数。

FX 图模式量化是 PyTorch 中的一个自动化量化框架，目前它是一个原型特性。它通过添加对函数的支持和自动化量化过程来改进急切模式量化，尽管用户可能需要重构模型以使其与 FX 图模式量化兼容（可以使用torch.fx符号化跟踪）。请注意，FX 图模式量化预计无法在任意模型上工作，因为模型可能无法符号化跟踪，我们将将其集成到诸如 torchvision 之类的领域库中，用户将能够使用 FX 图模式量化来量化类似于支持的领域库中的模型。对于任意模型，我们将提供一般指南，但要使其真正有效，用户可能需要熟悉torch.fx，尤其是在如何使模型符号化跟踪方面。

鼓励量化的新用户首先尝试 FX 图模式量化，如果它不起作用，用户可以尝试遵循使用 FX 图模式量化 的指南或回退到急切模式量化。

下表比较了急切模式量化和 FX 图模式量化之间的差异。

	急切模式量化	FX 图模式量化
发布状态	测试版	原型
算子融合	手动	自动
量化/反量化放置	手动	自动
量化模块	支持	支持
量化函数/Torch 操作	手动	自动
对自定义的支持	有限支持	完全支持
量化模式支持	训练后量化：静态、动态、仅权重 量化感知训练：静态	训练后量化：静态、动态、仅权重 量化感知训练：静态
输入/输出模型类型	torch.nn.Module	torch.nn.Module（可能需要一些重构才能使模型与 FX 图模式量化兼容）

后端/硬件支持

硬件	内核库	急切模式量化	FX 图模式量化	量化模式支持
服务器 CPU	fbgemm/onednn	支持		全部支持
移动 CPU	qnnpack/xnnpack
服务器 GPU	TensorRT（早期原型）	不支持，它需要一个图。	支持	静态量化

今天，PyTorch 支持以下后端以有效地运行量化算子：

• 支持 AVX2 或更高版本的 x86 CPU（没有 AVX2，某些操作的实现效率低下），通过fbgemm 和onednn 优化的x86（请参阅RFC 中的详细信息）。

• ARM CPU（通常在移动/嵌入式设备中找到），通过qnnpack。

• 通过TensorRT 通过fx2trt（将开源）对 NVidia GPU 的（早期原型）支持。

关于原生 CPU 后端的说明

我们使用相同的原生 PyTorch 量化算子公开x86 和qnnpack，因此我们需要额外的标志来区分它们。x86 和qnnpack 的相应实现根据 PyTorch 构建模式自动选择，尽管用户可以选择通过将torch.backends.quantization.engine 设置为x86 或qnnpack 来覆盖此设置。

在准备量化模型时，必须确保 qconfig 和用于量化计算的引擎与模型将在其上执行的后端匹配。qconfig 控制在量化过程中使用的观察器的类型。qengine 控制在为线性函数和卷积函数以及模块打包权重时是否使用x86 或qnnpack 特定的打包函数。例如：

x86 的默认设置

# set the qconfig for PTQ
# Note: the old 'fbgemm' is still available but 'x86' is the recommended default on x86 CPUs
qconfig = torch.ao.quantization.get_default_qconfig('x86')
# or, set the qconfig for QAT
qconfig = torch.ao.quantization.get_default_qat_qconfig('x86')
# set the qengine to control weight packing
torch.backends.quantized.engine = 'x86'

qnnpack 的默认设置

# set the qconfig for PTQ
qconfig = torch.ao.quantization.get_default_qconfig('qnnpack')
# or, set the qconfig for QAT
qconfig = torch.ao.quantization.get_default_qat_qconfig('qnnpack')
# set the qengine to control weight packing
torch.backends.quantized.engine = 'qnnpack'

算子支持

算子覆盖范围在动态量化和静态量化之间有所不同，并在下表中捕获。请注意，对于 FX 图模式量化，相应的函数也受支持。

	静态量化	动态量化
nn.Linear nn.Conv1d/2d/3d	Y Y	Y N
nn.LSTM nn.GRU	N N	Y Y
nn.RNNCell nn.GRUCell nn.LSTMCell	N N N	Y Y Y
nn.EmbeddingBag	Y（激活处于 fp32）	Y
nn.Embedding	Y	Y
nn.MultiheadAttention	不支持	不支持
激活	广泛支持	保持不变，计算保持在 fp32

注意：这将很快更新一些从原生 backend_config_dict 生成的信息。

量化 API 参考

量化 API 参考 包含量化 API 的文档，例如量化过程、量化张量操作以及支持的量化模块和函数。

量化后端配置

量化后端配置 包含有关如何为各种后端配置量化工作流程的文档。

量化精度调试

量化精度调试 包含有关如何调试量化精度的文档。

量化自定义

虽然提供了基于观察到的张量数据选择比例因子和偏差的观察器的默认实现，但开发人员可以提供他们自己的量化函数。可以有选择地将量化应用于模型的不同部分，或为模型的不同部分配置不同的量化方式。

我们还为conv1d()、conv2d()、conv3d() 和linear() 提供了每个通道量化的支持。

量化工作流程通过添加（例如，将观察器添加为.observer 子模块）或替换（例如，将nn.Conv2d 转换为nn.quantized.Conv2d）模型模块层次结构中的子模块来工作。这意味着模型在整个过程中保持常规的基于nn.Module 的实例，因此可以与 PyTorch 的其他 API 协同工作。

量化自定义模块 API

急切模式和 FX 图模式量化 API 都提供了一个钩子，供用户以自定义方式指定模块量化，并使用用户定义的观察和量化逻辑。用户需要指定：

1. 源 fp32 模块（存在于模型中）的 Python 类型。

2. 观察模块的 Python 类型（由用户提供）。此模块需要定义一个from_float 函数，该函数定义如何从原始 fp32 模块创建观察模块。

3. 量化模块的 Python 类型（由用户提供）。此模块需要定义一个from_observed 函数，该函数定义如何从观察模块创建量化模块。

4. 描述上述 (1)、(2)、(3) 的配置，传递给量化 API。

然后框架将执行以下操作：

1. 在prepare 模块交换期间，它将使用 (2) 中类的from_float 函数，将每个 (1) 中指定的类型的模块转换为 (2) 中指定的类型。

2. 在convert 模块交换期间，它将使用 (3) 中类的from_observed 函数，将每个 (2) 中指定的类型的模块转换为 (3) 中指定的类型。

目前，有一个要求是ObservedCustomModule 将具有单个张量输出，并且框架（而不是用户）将在该输出上添加一个观察器。观察器将作为自定义模块实例的属性存储在activation_post_process 键下。将来可能会放宽这些限制。

自定义 API 示例

import torch
import torch.ao.nn.quantized as nnq
from torch.ao.quantization import QConfigMapping
import torch.ao.quantization.quantize_fx

# original fp32 module to replace
class CustomModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 3)

    def forward(self, x):
        return self.linear(x)

# custom observed module, provided by user
class ObservedCustomModule(torch.nn.Module):
    def __init__(self, linear):
        super().__init__()
        self.linear = linear

    def forward(self, x):
        return self.linear(x)

    @classmethod
    def from_float(cls, float_module):
        assert hasattr(float_module, 'qconfig')
        observed = cls(float_module.linear)
        observed.qconfig = float_module.qconfig
        return observed

# custom quantized module, provided by user
class StaticQuantCustomModule(torch.nn.Module):
    def __init__(self, linear):
        super().__init__()
        self.linear = linear

    def forward(self, x):
        return self.linear(x)

    @classmethod
    def from_observed(cls, observed_module):
        assert hasattr(observed_module, 'qconfig')
        assert hasattr(observed_module, 'activation_post_process')
        observed_module.linear.activation_post_process = \
            observed_module.activation_post_process
        quantized = cls(nnq.Linear.from_float(observed_module.linear))
        return quantized

#
# example API call (Eager mode quantization)
#

m = torch.nn.Sequential(CustomModule()).eval()
prepare_custom_config_dict = {
    "float_to_observed_custom_module_class": {
        CustomModule: ObservedCustomModule
    }
}
convert_custom_config_dict = {
    "observed_to_quantized_custom_module_class": {
        ObservedCustomModule: StaticQuantCustomModule
    }
}
m.qconfig = torch.ao.quantization.default_qconfig
mp = torch.ao.quantization.prepare(
    m, prepare_custom_config_dict=prepare_custom_config_dict)
# calibration (not shown)
mq = torch.ao.quantization.convert(
    mp, convert_custom_config_dict=convert_custom_config_dict)
#
# example API call (FX graph mode quantization)
#
m = torch.nn.Sequential(CustomModule()).eval()
qconfig_mapping = QConfigMapping().set_global(torch.ao.quantization.default_qconfig)
prepare_custom_config_dict = {
    "float_to_observed_custom_module_class": {
        "static": {
            CustomModule: ObservedCustomModule,
        }
    }
}
convert_custom_config_dict = {
    "observed_to_quantized_custom_module_class": {
        "static": {
            ObservedCustomModule: StaticQuantCustomModule,
        }
    }
}
mp = torch.ao.quantization.quantize_fx.prepare_fx(
    m, qconfig_mapping, torch.randn(3,3), prepare_custom_config=prepare_custom_config_dict)
# calibration (not shown)
mq = torch.ao.quantization.quantize_fx.convert_fx(
    mp, convert_custom_config=convert_custom_config_dict)

最佳实践

1. 如果您使用的是 x86 后端，我们需要使用 7 位而不是 8 位。请确保您减少了 quant\_min 和 quant\_max 的范围，例如，如果 dtype 是 torch.quint8，请确保将自定义的 quant_min 设置为 0，并将 quant_max 设置为 127（255 / 2）；如果 dtype 是 torch.qint8，请确保将自定义的 quant_min 设置为 -64（-128 / 2），并将 quant_max 设置为 63（127 / 2）。如果您调用 torch.ao.quantization.get_default_qconfig(backend) 或 torch.ao.quantization.get_default_qat_qconfig(backend) 函数来获取 x86 或 qnnpack 后端的默认 qconfig，我们已经正确设置了这些值。

2. 如果选择了 onednn 后端，则默认的 qconfig 映射 torch.ao.quantization.get_default_qconfig_mapping('onednn') 和默认的 qconfig torch.ao.quantization.get_default_qconfig('onednn') 将为激活使用 8 位。建议在支持矢量神经网络指令 (VNNI) 的 CPU 上使用它。否则，将激活观察器的 reduce_range 设置为 True，以便在不支持 VNNI 的 CPU 上获得更好的精度。

常见问题

1. 如何在 GPU 上进行量化推理？

   我们还没有官方的 GPU 支持，但这是一个积极开发的领域，您可以在 这里 找到更多信息。

2. 在哪里可以获得我量化模型的 ONNX 支持？

   如果您在导出模型时遇到错误（使用 torch.onnx 下的 API），您可以打开 PyTorch 存储库中的问题。在问题标题前加上 [ONNX] 并将问题标记为 module: onnx。

   如果您遇到 ONNX Runtime 的问题，请在 GitHub - microsoft/onnxruntime 中打开问题。

3. 如何将量化与 LSTM 一起使用？

   LSTM 通过我们在 Eager 模式和 FX 图模式量化中的自定义模块 API 得到支持。示例可以在以下位置找到：Eager 模式：pytorch/test_quantized_op.py TestQuantizedOps.test_custom_module_lstm FX 图模式：pytorch/test_quantize_fx.py TestQuantizeFx.test_static_lstm

常见错误

将非量化张量传递到量化内核

如果您看到类似以下的错误：

RuntimeError: Could not run 'quantized::some_operator' with arguments from the 'CPU' backend...

这意味着您正在尝试将非量化张量传递到量化内核。一种常见的解决方法是使用 torch.ao.quantization.QuantStub 来量化张量。这需要在 Eager 模式量化中手动完成。一个端到端示例

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = torch.ao.quantization.QuantStub()
        self.conv = torch.nn.Conv2d(1, 1, 1)

    def forward(self, x):
        # during the convert step, this will be replaced with a
        # `quantize_per_tensor` call
        x = self.quant(x)
        x = self.conv(x)
        return x

将量化张量传递到非量化内核

如果您看到类似以下的错误：

RuntimeError: Could not run 'aten::thnn_conv2d_forward' with arguments from the 'QuantizedCPU' backend.

这意味着您正在尝试将量化张量传递到非量化内核。一种常见的解决方法是使用 torch.ao.quantization.DeQuantStub 来反量化张量。这需要在 Eager 模式量化中手动完成。一个端到端示例

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = torch.ao.quantization.QuantStub()
        self.conv1 = torch.nn.Conv2d(1, 1, 1)
        # this module will not be quantized (see `qconfig = None` logic below)
        self.conv2 = torch.nn.Conv2d(1, 1, 1)
        self.dequant = torch.ao.quantization.DeQuantStub()

    def forward(self, x):
        # during the convert step, this will be replaced with a
        # `quantize_per_tensor` call
        x = self.quant(x)
        x = self.conv1(x)
        # during the convert step, this will be replaced with a
        # `dequantize` call
        x = self.dequant(x)
        x = self.conv2(x)
        return x

m = M()
m.qconfig = some_qconfig
# turn off quantization for conv2
m.conv2.qconfig = None

保存和加载量化模型

在对量化模型调用 torch.load 时，如果您看到类似以下的错误：

AttributeError: 'LinearPackedParams' object has no attribute '_modules'

这是因为使用 torch.save 和 torch.load 直接保存和加载量化模型不受支持。要保存/加载量化模型，可以使用以下方法

1. 保存/加载量化模型的 state_dict

一个例子

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(5, 5)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)
        return x

m = M().eval()
prepare_orig = prepare_fx(m, {'' : default_qconfig})
prepare_orig(torch.rand(5, 5))
quantized_orig = convert_fx(prepare_orig)

# Save/load using state_dict
b = io.BytesIO()
torch.save(quantized_orig.state_dict(), b)

m2 = M().eval()
prepared = prepare_fx(m2, {'' : default_qconfig})
quantized = convert_fx(prepared)
b.seek(0)
quantized.load_state_dict(torch.load(b))

2. 使用 torch.jit.save 和 torch.jit.load 保存/加载脚本化的量化模型

一个例子

# Note: using the same model M from previous example
m = M().eval()
prepare_orig = prepare_fx(m, {'' : default_qconfig})
prepare_orig(torch.rand(5, 5))
quantized_orig = convert_fx(prepare_orig)

# save/load using scripted model
scripted = torch.jit.script(quantized_orig)
b = io.BytesIO()
torch.jit.save(scripted, b)
b.seek(0)
scripted_quantized = torch.jit.load(b)

使用 FX 图模式量化时出现符号跟踪错误

符号可跟踪性是 (原型 - 维护模式) FX 图模式量化 的一项要求，因此，如果您将一个不可符号跟踪的 PyTorch 模型传递给 torch.ao.quantization.prepare_fx 或 torch.ao.quantization.prepare_qat_fx，我们可能会看到如下错误：

torch.fx.proxy.TraceError: symbolically traced variables cannot be used as inputs to control flow

请查看 符号跟踪的局限性 并使用 - 使用 FX 图模式量化的用户指南 来解决此问题。