量化¶

警告

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

量化简介¶

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

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

在较低层次上，PyTorch 提供了一种表示量化张量的方法，并使用它们执行操作。它们可以用来直接构建模型，这些模型可以在较低精度下执行全部或部分计算。提供了更高层次的 API，这些 API 包含了将 FP32 模型转换为较低精度且精度损失最小的典型工作流程。

量化 API 概述¶

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

Eager Mode Quantization 是一个测试版功能。用户需要手动进行融合，并指定量化和反量化发生的位置，同时它仅支持模块而不支持功能。

FX Graph Mode Quantization 是 PyTorch 中的一个自动化量化工作流程，目前它是一个原型功能，自 PyTorch 2 Export Quantization 以来一直处于维护模式。它通过增加对功能的支持并自动化量化过程，改进了 Eager Mode Quantization，尽管人们可能需要重构模型以使其与 FX Graph Mode Quantization 兼容（能够通过 torch.fx 进行符号追踪）。需要注意的是，FX Graph Mode Quantization 并不期望在任意模型上工作，因为模型可能无法进行符号追踪，我们将会将其集成到像 torchvision 这样的领域库中，用户将能够使用 FX Graph Mode Quantization 量化类似于支持领域库中的模型。对于任意模型，我们将提供一般性指导，但要实际使其工作，用户可能需要熟悉 torch.fx，特别是如何使模型能够进行符号追踪。

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

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

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

	急切模式量化	FX 图表模式量化	PyTorch 2 导出量化
发布状态	贝塔	原型 (维护)	原型
操作符融合	手册	自动	自动
量化/反量化放置	手册	自动	自动
量化模块	支持	支持	支持
量化功能/Torch 操作	手册	自动	支持
支持自定义	有限支持	完全支持	完全支持
量化模式支持	训练后量化：静态、动态、仅权重量化感知训练: 静态	训练后量化：静态、动态、仅权重量化感知训练: 静态	由后端特定量化器定义
输入/输出模型类型	`torch.nn.Module`	`torch.nn.Module` （可能需要一些重构以使模型与FX图模式量化兼容）	`torch.fx.GraphModule` （由 `torch.export`捕获）

支持三种类型的量化：

动态量化（权重被量化，激活值以浮点数读取/存储，并在计算时量化）
静态量化（权重量化，激活量化，训练后需要校准）
静态量化感知训练（权重量化，激活值量化，训练期间建模量化数值）

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

动态和静态量化之间的操作符覆盖范围有所不同，具体内容见下表。

	静态量化	动态量化
nn.Linear nn.Conv1d/2d/3d	Y Y	Y N
nn.LSTM nn.GRU	Y (through custom modules) N	Y Y
nn.RNNCell nn.GRUCell nn.LSTMCell	N N N	Y Y Y
nn.EmbeddingBag	Y（激活值为fp32格式）	Y
nn.嵌入	Y	Y
nn.多头注意力机制	Y（通过自定义模块）	不支持
激活函数	广泛支持	未改变，计算保持在 fp32

急切模式量化¶

有关量化流程的总体介绍，包括不同类型的量化，请参阅量化流程概述。

训练后动态量化¶

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

图表：

# 原始模型
# 所有张量和计算都是浮点数
previous_layer_fp32 -- linear_fp32 -- activation_fp32 -- next_layer_fp32
                 /
linear_weight_fp32

# 动态量化模型
# 线性和LSTM权重为int8
previous_layer_fp32 -- linear_int8_w_fp32_inp -- activation_fp32 -- next_layer_fp32
                     /
   linear_weight_int8

PTDQ API 示例：

import torch

# 定义一个浮点模型
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

# 创建一个模型实例
model_fp32 = M()
# 创建一个量化模型实例
model_int8 = torch.ao.quantization.quantize_dynamic(
    model_fp32,  # 原始模型
    {torch.nn.Linear},  # 要动态量化的层集合
    dtype=torch.qint8)  # 量化权重的目标数据类型

# 运行模型
input_fp32 = torch.randn(4, 4, 4, 4)
res = model_int8(input_fp32)

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

训练后静态量化¶

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

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

图表：

# 原始模型
# 所有张量和计算都是浮点数
previous_layer_fp32 -- linear_fp32 -- activation_fp32 -- next_layer_fp32
                    /
    linear_weight_fp32

# 静态量化模型
# 权重和激活值都是int8
previous_layer_int8 -- linear_with_activation_int8 -- next_layer_int8
                    /
  linear_weight_int8

PTSQ API 示例：

import torch

# 定义一个浮点模型，其中某些层可以静态量化
class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # QuantStub将张量从浮点数转换为量化
        self.quant = torch.ao.quantization.QuantStub()
        self.conv = torch.nn.Conv2d(1, 1, 1)
        self.relu = torch.nn.ReLU()
        # DeQuantStub将张量从量化转换为浮点数
        self.dequant = torch.ao.quantization.DeQuantStub()

    def forward(self, x):
        # 手动指定张量将在量化模型中从浮点数转换为量化的位置
        x = self.quant(x)
        x = self.conv(x)
        x = self.relu(x)
        # 手动指定张量将在量化模型中从量化转换为浮点数的位置
        x = self.dequant(x)
        return x

# 创建一个模型实例
model_fp32 = M()

# 模型必须设置为eval模式，以便静态量化逻辑正常工作
model_fp32.eval()

# 附加一个全局qconfig，其中包含有关要附加的观察者的信息。使用'x86'进行服务器推理，使用'qnnpack'进行移动推理。其他量化配置，如选择对称或非对称量化以及MinMax或L2Norm校准技术，可以在此处指定。
# 注意：旧的'fbgemm'仍然可用，但'x86'是服务器推理的推荐默认值。
# model_fp32.qconfig = torch.ao.quantization.get_default_qconfig('fbgemm')
model_fp32.qconfig = torch.ao.quantization.get_default_qconfig('x86')

# 将激活融合到前面的层中，如果适用的话。这需要根据模型架构手动完成。常见的融合包括`conv + relu`和`conv + batchnorm + relu`
model_fp32_fused = torch.ao.quantization.fuse_modules(model_fp32, [['conv', 'relu']])

# 准备模型进行静态量化。这会在模型中插入观察者，这些观察者将在校准时观察激活张量。
model_fp32_prepared = torch.ao.quantization.prepare(model_fp32_fused)

# 校准准备好的模型以确定激活的量化参数。在实际应用中，校准将使用代表性数据集完成。
input_fp32 = torch.randn(4, 1, 4, 4)
model_fp32_prepared(input_fp32)

# 将观察到的模型转换为量化模型。这会执行以下操作：量化权重，计算并存储与每个激活张量一起使用的比例和偏差值，并用量化实现替换关键操作符。
model_int8 = torch.ao.quantization.convert(model_fp32_prepared)

# 运行模型，相关的计算将在int8中进行
res = model_int8(input_fp32)

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

静态量化的量化感知训练¶

量化感知训练 (QAT) 在训练过程中模拟量化的效果，与其他量化方法相比，可以获得更高的准确性。我们可以对静态、动态或仅权重量化进行 QAT。在训练过程中，所有计算都以浮点数进行，通过 fake_quant 模块模拟量化的效果，通过钳位和舍入来模拟 INT8 的效果。模型转换后，权重和激活值被量化，并且尽可能将激活值融合到前一层中。它通常与卷积神经网络 (CNN) 一起使用，与静态量化相比，可以获得更高的准确性。

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

图表：

# 原始模型
# 所有张量和计算都是浮点数
previous_layer_fp32 -- linear_fp32 -- activation_fp32 -- next_layer_fp32
                      /
    linear_weight_fp32

# 带有fake_quants的模型，用于在训练期间模拟量化数值
previous_layer_fp32 -- fq -- linear_fp32 -- activation_fp32 -- fq -- next_layer_fp32
                           /
   linear_weight_fp32 -- fq

# 量化模型
# 权重和激活值为int8
previous_layer_int8 -- linear_with_activation_int8 -- next_layer_int8
                     /
   linear_weight_int8

QAT API 示例：

import torch

# 定义一个浮点模型，其中某些层可以从QAT中受益
class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # QuantStub将张量从浮点转换为量化
        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将张量从量化转换为浮点
        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

# 创建一个模型实例
model_fp32 = M()

# 模型必须设置为eval模式以进行融合
model_fp32.eval()

# 附加一个全局qconfig，其中包含关于附加哪种观察者的信息。使用'x86'用于服务器推理，'qnnpack'用于移动推理。其他量化配置，如选择对称或非对称量化以及MinMax或L2Norm校准技术，可以在此处指定。
# 注意：旧的'fbgemm'仍然可用，但'x86'是服务器推理的推荐默认值。
# model_fp32.qconfig = torch.ao.quantization.get_default_qconfig('fbgemm')
model_fp32.qconfig = torch.ao.quantization.get_default_qat_qconfig('x86')

# 将激活融合到前面的层中，如果适用
# 这需要根据模型架构手动完成
model_fp32_fused = torch.ao.quantization.fuse_modules(model_fp32,
    [['conv', 'bn', 'relu']])

# 为QAT准备模型。这将在模型中插入观察者和fake_quants
# 模型需要设置为train模式以使QAT逻辑工作
# 在校准期间观察权重和激活张量的模型。
model_fp32_prepared = torch.ao.quantization.prepare_qat(model_fp32_fused.train())

# 运行训练循环（未显示）
training_loop(model_fp32_prepared)

# 将观察到的模型转换为量化模型。这做了几件事：
# 量化权重，计算并存储与每个激活张量一起使用的比例和偏差值，在适当的地方融合模块，
# 并将关键操作符替换为量化实现。
model_fp32_prepared.eval()
model_int8 = torch.ao.quantization.convert(model_fp32_prepared)

# 运行模型，相关的计算将在int8中进行
res = model_int8(input_fp32)

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

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

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

将任何需要输出重量化（因此具有额外参数）的操作从功能形式转换为模块形式（例如，使用torch.nn.ReLU而不是torch.nn.functional.relu）。
通过在子模块上分配.qconfig属性或通过指定qconfig_mapping来指定模型中需要量化的部分。例如，设置model.conv1.qconfig = None意味着model.conv层将不会被量化，而设置model.linear1.qconfig = custom_qconfig意味着model.linear1的量化设置将使用custom_qconfig而不是全局qconfig。

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

指定激活值的量化和反量化位置。这是通过使用 QuantStub 和 DeQuantStub 模块来完成的。
使用 FloatFunctional 将需要特殊处理的量化张量操作包装到模块中。例如，像 add 和 cat 这样的操作需要特殊处理以确定输出量化参数。
融合模块：将操作/模块组合成一个单一模块，以获得更高的准确性和性能。这是通过使用 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()

#
# 训练后动态/仅权重量化
#

# 我们需要深度复制，如果我们仍然想在量化后保持model_fp不变，因为量化API会改变输入模型
model_to_quantize = copy.deepcopy(model_fp)
model_to_quantize.eval()
qconfig_mapping = QConfigMapping().set_global(torch.ao.quantization.default_dynamic_qconfig)
# 需要一个或多个示例输入的元组来跟踪模型
example_inputs = (input_fp32)
# 准备
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, example_inputs)
# 当我们只有动态/仅权重量化时，不需要校准
# 量化
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# 训练后静态量化
#

model_to_quantize = copy.deepcopy(model_fp)
qconfig_mapping = get_default_qconfig_mapping("qnnpack")
model_to_quantize.eval()
# 准备
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, example_inputs)
# 校准（未显示）
# 量化
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# 静态量化的量化感知训练
#

model_to_quantize = copy.deepcopy(model_fp)
qconfig_mapping = get_default_qat_qconfig_mapping("qnnpack")
model_to_quantize .train()
# 准备
model_prepared = quantize_fx.prepare_qat_fx(model_to_quantize, qconfig_mapping, example_inputs)
# 训练循环（未显示）
# 量化
model_quantized = quantize_fx.convert_fx(model_prepared)

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

请按照以下教程了解更多关于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)

# 初始化一个浮点模型
float_model = M().eval()

# 定义校准函数
def calibrate(model, data_loader):
    model.eval()
    with torch.no_grad():
        for image, target in data_loader:
            model(image)

# 步骤1. 程序捕获
# 注意：此API将在未来更新为torch.export API，但捕获的结果应该基本保持不变
m = capture_pre_autograd_graph(m, *example_inputs)
# 我们得到一个包含aten操作的模型

# 步骤2. 量化
# 后端开发者将编写他们自己的Quantizer并公开方法，以允许用户表达他们
# 希望如何对模型进行量化
quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config())
# 或者为量化感知训练准备prepare_qat_pt2e
m = prepare_pt2e(m, quantizer)

# 运行校准
# calibrate(m, sample_inference_data)
m = convert_pt2e(m)

# 步骤3. 降低
# 降低到目标后端

请按照以下教程开始使用 PyTorch 2 导出量化：

建模用户：

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

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

量化栈¶

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

量化模型¶

量化张量¶

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

PyTorch 支持逐张量和逐通道的对称和非对称量化。逐张量意味着张量中的所有值都使用相同的量化参数进行量化。逐通道意味着对于每个维度（通常是张量的通道维度），张量中的值使用不同的量化参数进行量化。这允许在将张量转换为量化值时减少误差，因为异常值只会影响其所在的通道，而不是整个张量。

映射是通过将浮点张量转换为

$_images/math-quantizer-equation.png$

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

以下是量化张量的几个关键属性：

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 将具有量化参数
  - 比例 (浮点数)
  - zero_point (整数)
- torch.per_channel_affine 将具有的量化参数为
  - per_channel_scales (浮点数列表)
  - per_channel_zero_points (整数列表)
  - 轴 (int)

量化与反量化¶

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

量化 (浮点数 -> 量化值)
- 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)
- 转换为(torch.float16)
反量化（量化值 -> 浮点数）
- quantized_tensor.dequantize() - 对一个 torch.float16 张量调用 dequantize 将把张量转换回 torch.float
- torch.dequantize(x)

量化操作符/模块¶

量化算子是那些以量化张量作为输入，并输出量化张量的算子。
量化模块是执行量化操作的 PyTorch 模块。它们通常为线性操作和卷积等加权操作定义。

量化引擎¶

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

量化流程¶

观察者和伪量化¶

观察者是用于以下目的的 PyTorch 模块：
- 收集张量统计信息，如通过观察者的张量的最小值和最大值
- 并根据收集的张量统计信息计算量化参数
FakeQuantize 是 PyTorch 模块，用于：
- 模拟量化（执行量化/反量化）网络中的张量
- 它可以根据从观察者收集的统计数据计算量化参数，也可以学习量化参数

QConfig¶

QConfig 是一个包含 Observer 或 FakeQuantize 模块类的 namedtuple，可以通过 qscheme、dtype 等进行配置。它用于配置如何观察一个操作符
- 操作符/模块的量化配置
  - 不同类型的 Observer/FakeQuantize
  - 数据类型
  - qscheme
  - quant_min/quant_max: 可用于模拟低精度张量
- 目前支持激活和权重的配置
- 我们根据为给定操作符或模块配置的qconfig插入输入/权重/输出观察器

通用量化流程¶

一般来说，流程如下

准备
- 根据用户指定的qconfig插入Observer/FakeQuantize模块
校准/训练（取决于训练后量化或量化感知训练）
- 允许观察者收集统计数据或FakeQuantize模块学习量化参数
转换
- 将校准/训练后的模型转换为量化模型

量化有不同的模式，它们可以分为两种方式：

关于我们应用量化流程的地方，我们有：

训练后量化（训练后应用量化，量化参数基于样本校准数据计算）
量化感知训练（在训练过程中模拟量化，以便可以使用训练数据与模型一起学习量化参数）

在量化算子的方式上，我们可以有：

仅权重量化（仅权重是静态量化的）
动态量化（权重是静态量化的，激活是动态量化的）
静态量化（权重和激活值都是静态量化的）

我们可以在同一个量化流程中混合使用不同的量化操作符方式。例如，我们可以有包含静态和动态量化操作符的后训练量化。

量化支持矩阵¶

量化模式支持¶

	量化模式		数据集要求	最适合用于	准确性	注释
训练后量化	动态/仅权重量化	激活动态量化（fp16， int8）或不量化，权重静态量化（fp16， int8， in4）	无	LSTM, MLP, 嵌入, Transformer	好	易于使用，接近静态量化当性能是计算或内存受限于权重
训练后量化	静态量化	激活和权重静态量化（int8）	校准数据集	卷积神经网络	好	提供最佳性能，可能对准确性有较大影响，适合仅支持int8计算的硬件
量化感知训练	动态量化	激活和权重是伪量化的	微调数据集	MLP, 嵌入	最好	目前仅提供有限支持
量化感知训练	静态量化	激活和权重是伪量化的	微调数据集	卷积神经网络, 多层感知机, 嵌入	最好	通常用于当静态量化导致精度不佳时，用于缩小精度差距

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

量化流程支持¶

PyTorch 提供了两种量化模式：Eager 模式量化和 FX 图模式量化。

Eager Mode Quantization 是一个测试版功能。用户需要手动进行融合，并指定量化和反量化发生的位置，同时它仅支持模块而不支持功能。

FX Graph Mode Quantization 是 PyTorch 中的一个自动化量化框架，目前它是一个原型功能。它通过增加对功能的支持并自动化量化过程，改进了 Eager Mode Quantization，尽管人们可能需要重构模型以使其与 FX Graph Mode Quantization 兼容（能够通过 torch.fx 进行符号追踪）。需要注意的是，FX Graph Mode Quantization 并不期望在任意模型上都能工作，因为模型可能无法进行符号追踪，我们将会将其集成到像 torchvision 这样的领域库中，用户将能够使用 FX Graph Mode Quantization 量化类似于支持领域库中的模型。对于任意模型，我们将提供一般性指导，但要实际使其工作，用户可能需要熟悉 torch.fx，特别是如何使模型能够进行符号追踪。

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

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

	急切模式量化	FX 图表模式量化
发布状态	贝塔	原型
操作符融合	手册	自动
量化/反量化放置	手册	自动
量化模块	支持	支持
量化功能/Torch 操作	手册	自动
支持自定义	有限支持	完全支持
量化模式支持	训练后量化：静态、动态、仅权重量化感知训练: 静态	训练后量化：静态、动态、仅权重量化感知训练: 静态
输入/输出模型类型	`torch.nn.Module`	`torch.nn.Module` （可能需要一些重构以使模型与FX图模式量化兼容）

后端/硬件支持¶

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

今天，PyTorch 支持以下后端以高效运行量化操作符：

支持AVX2或更高版本的x86 CPU（没有AVX2时，某些操作的实现效率较低），通过x86由fbgemm和onednn优化（详见RFC）
ARM CPU（通常用于移动/嵌入式设备），通过 qnnpack
（早期原型）通过 TensorRT 支持 NVidia GPU，使用 fx2trt（即将开源）

原生CPU后端的注意事项¶

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

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

x86 的默认设置：

# 设置PTQ的qconfig
# 注意：旧的'fbgemm'仍然可用，但在x86 CPU上'x86'是推荐的默认值
qconfig = torch.ao.quantization.get_default_qconfig('x86')
# 或者，设置QAT的qconfig
qconfig = torch.ao.quantization.get_default_qat_qconfig('x86')
# 设置qengine以控制权重打包
torch.backends.quantized.engine = 'x86'

qnnpack 的默认设置：

# 设置PTQ的qconfig
qconfig = torch.ao.quantization.get_default_qconfig('qnnpack')
# 或者，设置QAT的qconfig
qconfig = torch.ao.quantization.get_default_qat_qconfig('qnnpack')
# 设置qengine以控制权重打包
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.嵌入	Y	Y
nn.多头注意力机制	不支持	不支持
激活函数	广泛支持	未改变，计算保持在 fp32

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

量化 API 参考¶

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

量化后端配置¶

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

量化精度调试¶

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

量化自定义¶

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

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

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

量化自定义模块API¶

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

源 fp32 模块（存在于模型中）的 Python 类型
观察模块的Python类型（由用户提供）。该模块需要定义一个from_float函数，该函数定义了如何从原始的fp32模块创建观察模块。
量化模块的Python类型（由用户提供）。该模块需要定义一个from_observed函数，该函数定义了如何从观察到的模块创建量化模块。
描述上述(1)、(2)、(3)的配置，传递给量化API。

框架将会执行以下操作：

在准备模块交换期间，它将把（1）中指定的每个模块类型转换为（2）中指定的类型，使用（2）中类的from_float函数。
在convert模块交换期间，它将把(2)中指定的每个模块类型转换为(3)中指定的类型，使用(3)中类的from_observed函数。

目前，有一个要求是ObservedCustomModule将有一个单一的Tensor输出，并且框架（不是用户）将在这个输出上添加一个观察器。该观察器将作为自定义模块实例的一个属性存储在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

# 原始的 fp32 模块，用于替换
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)

# 自定义的观察模块，由用户提供
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

# 自定义的量化模块，由用户提供
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

#
# 示例 API 调用（Eager 模式量化）
#

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)
# 校准（未显示）
mq = torch.ao.quantization.convert(
    mp, convert_custom_config_dict=convert_custom_config_dict)
#
# 示例 API 调用（FX 图模式量化）
#
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)
# 校准（未显示）
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) 函数来获取默认的 qconfig，我们已经正确设置了这些值。

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

常见问题¶

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

我们尚未提供官方的GPU支持，但这是我们正在积极开发的一个领域，您可以在此处找到更多信息 here
我可以在哪里获得量化模型的ONNX支持？

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

如果您在使用ONNX Runtime时遇到问题，请在GitHub - microsoft/onnxruntime上提交问题。
如何将量化应用于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: 无法运行 'quantized::some_operator' 与 来自 'CPU' 后端的参数...

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

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):
        # 在转换步骤中，这将替换为
        # `quantize_per_tensor` 调用
        x = self.quant(x)
        x = self.conv(x)
        return x

将量化张量传递给非量化内核¶

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

RuntimeError: 无法 运行 'aten::thnn_conv2d_forward' 使用 来自 'QuantizedCPU' 后端 的参数.

这意味着你正在尝试将一个量化张量传递给一个非量化内核。一个常见的解决方法是使用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)
        # 这个模块将不会被量化（见下面的 `qconfig = None` 逻辑）
        self.conv2 = torch.nn.Conv2d(1, 1, 1)
        self.dequant = torch.ao.quantization.DeQuantStub()

    def forward(self, x):
        # 在转换步骤中，这将被替换为一个 `quantize_per_tensor` 调用
        x = self.quant(x)
        x = self.conv1(x)
        # 在转换步骤中，这将被替换为一个 `dequantize` 调用
        x = self.dequant(x)
        x = self.conv2(x)
        return x

m = M()
m.qconfig = some_qconfig
# 关闭 conv2 的量化
m.conv2.qconfig = None

保存和加载量化模型¶

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

AttributeError: 'LinearPackedParams' 对象 没有 属性 '_modules'

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

保存/加载量化模型的状态字典

一个例子：

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)

# 使用 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))

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

一个例子：

# 注意：使用与前一个示例相同的模型 M
m = M().eval()
prepare_orig = prepare_fx(m, {'' : default_qconfig})
prepare_orig(torch.rand(5, 5))
quantized_orig = convert_fx(prepare_orig)

# 使用脚本模型保存/加载
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: 符号化跟踪的变量不能用作控制流的输入

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