量化

此文件正在迁移至 torch/ao/quantization，在此期间保留于此以确保兼容性。 如果要添加新条目/功能，请将其添加到 torch/ao/quantization/fx/ 下的相应文件中，并在此处添加导入语句。

警告

量化功能目前处于测试阶段，可能会有所变更。

量化简介

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

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

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

量化 API 概述

急切模式量化

关于量化流程的一般介绍，包括不同类型的量化，请参阅通用量化流程。

训练后动态量化

这是最简单的量化形式，其中权重在提前进行量化，而激活值则在推理过程中动态量化。这种情况适用于模型的执行时间主要消耗在从内存加载权重而不是计算矩阵乘法上。对于小批量处理的LSTM和Transformer类型的模型来说，这种情况是成立的。

Diagram:

# 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

API 示例：

import torch

# define a floating point model
class M(torch.nn.Module):
    def __init__(self):
        super(M, self).__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.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 static）对模型的权重和激活值进行量化。它会在可能的情况下将激活值与前面的层融合。它需要使用具有代表性的数据集进行校准，以确定激活值的最佳量化参数。当内存带宽和计算节省都很重要时，通常会使用训练后静态量化，卷积神经网络（CNNs）是其典型应用场景。

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

Diagram:

# 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

API 示例：

import torch

# define a floating point model where some layers could be statically quantized
class M(torch.nn.Module):
    def __init__(self):
        super(M, self).__init__()
        # QuantStub converts tensors from floating point to quantized
        self.quant = torch.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.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 'fbgemm' for server inference and
# 'qnnpack' for mobile inference. Other quantization configurations such
# as selecting symmetric or assymetric quantization and MinMax or L2Norm
# calibration techniques can be specified here.
model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 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.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.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.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 一起使用，相比静态量化能提供更高的精度。

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

Diagram:

# 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

API 示例：

import torch

# define a floating point model where some layers could benefit from QAT
class M(torch.nn.Module):
    def __init__(self):
        super(M, self).__init__()
        # QuantStub converts tensors from floating point to quantized
        self.quant = torch.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.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 train mode for QAT logic to work
model_fp32.train()

# attach a global qconfig, which contains information about what kind
# of observers to attach. Use 'fbgemm' for server inference and
# 'qnnpack' for mobile inference. Other quantization configurations such
# as selecting symmetric or assymetric quantization and MinMax or L2Norm
# calibration techniques can be specified here.
model_fp32.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')

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

# Prepare the model for QAT. This inserts observers and fake_quants in
# the model that will observe weight and activation tensors during calibration.
model_fp32_prepared = torch.quantization.prepare_qat(model_fp32_fused)

# 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.quantization.convert(model_fp32_prepared)

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

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

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

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

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

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

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

1. 指定在哪里对激活值进行量化和反量化。这是通过 QuantStub 和 DeQuantStub 模块完成的。

2. 使用 torch.nn.quantized.FloatFunctional 将需要对量化进行特殊处理的张量操作封装成模块。例如像 add 和 cat 这样的操作，它们需要特殊的处理来确定输出量化的参数。

3. 融合模块：将操作/模块组合成一个模块，以获得更高的准确性和性能。这是通过 torch.quantization.fuse_modules() API 实现的，该 API 接收要融合的模块列表。我们目前支持以下融合方式： [卷积, ReLU], [卷积, 批归一化], [卷积, 批归一化, ReLU], [线性, ReLU]

(原型) FX 图模式量化

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

API 示例：

import torch.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_dict = {"": torch.quantization.default_dynamic_qconfig}
# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_dict)
# no calibration needed when we only have dynamici/weight_only quantization
# quantize
model_quantized = quantize_fx.convert_fx(model_prepared)

#
# post training static quantization
#

model_to_quantize = copy.deepcopy(model_fp)
qconfig_dict = {"": torch.quantization.get_default_qconfig('qnnpack')}
model_to_quantize.eval()
# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_dict)
# 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_dict = {"": torch.quantization.get_default_qat_qconfig('qnnpack')}
model_to_quantize.train()
# prepare
model_prepared = quantize_fx.prepare_qat_fx(model_to_quantize, qconfig_dict)
# 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 Graph Mode Quantization）的更多信息：

• 使用 FX 图模式量化指南

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

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

量化堆栈

量化是将浮点模型转换为量化模型的过程。因此，从整体来看，量化堆栈可以分为两个部分：1). 用于量化模型的构建模块或抽象；2). 用于将浮点模型转换为量化模型的量化流程的构建模块或抽象

量化模型

量化张量

在 PyTorch 中进行量化，我们需要能够在张量中表示量化数据。量化张量允许存储量化数据（以 int8/uint8/int32 形式表示），并包含如比例因子和零点等量化参数。量化张量支持许多有用的操作，使量化算术变得简单，并且还允许以量化格式对数据进行序列化。

PyTorch 支持每张量和每通道的对称和非对称量化。每张量表示张量内的所有值都使用相同的量化参数以相同的方式进行量化。每通道表示对于每个维度，通常是张量的通道维度，张量中的值会使用不同的量化参数进行量化。这使得在将张量转换为量化值时可以减少误差，因为异常值只会对其所在的通道产生影响，而不是整个张量。

通过将浮点张量转换来执行映射。

请注意，我们确保浮点数中的零在量化后能够无误差地表示，从而保证填充等操作不会引入额外的量化误差。

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

• QScheme (torch.qscheme)：一个枚举类型，指定我们量化 Tensor 的方式

  • 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 (int)

  • torch.per_tensor_affine 的量化参数为

    • per_channel_scales (浮点数列表)

    • per_channel_zero_points (整数列表)

    • 轴（int）

量化和反量化

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

• 量化 (float -> quantized)

  • 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 类型的 Tensor 调用 dequantize 将会将其转换回 torch.float

  • torch.dequantize(x)

量化操作符/模块

• 量化操作符是将量化张量作为输入，并输出量化张量的操作符。

• 量化模块是执行量化操作的 PyTorch 模块。它们通常用于带权重的操作，如线性层和卷积层。

量化引擎

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

量化流程

观察器和伪量化

• Observer 是用于以下目的的 PyTorch 模块：

  • 收集张量统计数据，例如通过观察器的 Tensor 的最小值和最大值

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

• FakeQuantize 是 PyTorch 模块，用于：

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

  • 它可以基于观察器收集的统计信息计算量化参数，或者也可以学习量化参数

QConfig

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

  • 操作符/模块的量化配置

    • 不同类型的Observer/FakeQuantize

    • 数据类型(dtype)

    • 量化方案

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

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

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

通用量化流程

通常，流程如下

• 准备

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

• 校准/训练（根据训练后量化或量化感知训练而定）

  • 允许 Observers 收集统计信息或 FakeQuantize 模块学习量化参数

• 转换

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

量化有多种模式，它们可以通过两种方式进行分类：

在量化流程的应用位置方面，我们有：

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

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

在如何对操作符进行量化方面，我们可以有：

• 仅权重量化（仅权重进行静态量化）

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

• 静态量化（权重和激活值均被静态量化）

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

量化支持矩阵

量化模式支持

	量化 模式		数据集 需求	适用于	准确率	注释
训练后量化	动态/仅权重量化	激活 动态 量化（fp16，int8）或未量化 量化，权重 静态量化 （fp16，int8，in4）	请提供需要翻译的单词列表。	LSTM，MLP， Embedding， Transformer	好	易于使用， 接近静态量化， 当性能受限于 计算或内存时， 由于权重
	静态量化	激活和 权重静态量化（int8）	校准 数据集	CNN	好	提供最佳性能，可能对精度产生较大影响，适用于仅支持 int8 计算的硬件
量化感知训练	动态量化	激活和 权重是伪量化	微调 数据集	多层感知机，嵌入	最佳	目前支持有限
	静态量化	激活和 权重是伪量化	微调 数据集	CNN，MLP， Embedding	最佳	通常用于静态 量化导致精度 下降的情况，用以 弥补精度差距

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

量化流程支持

PyTorch 提供了两种量化模式：Eager Mode Quantization（即时模式量化）和 FX Graph Mode Quantization（FX 图模式量化）。

Eager Mode 量化是一个测试功能。用户需要手动执行融合操作并指定量化和反量化的位置，同时它仅支持模块，不支持函数式操作。

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

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

下表比较了 Eager Mode Quantization（急切模式量化）和 FX Graph Mode Quantization（FX 图模式量化）之间的差异：

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

后端/硬件支持

硬件	内核库	Eager Mode 量化	FX 图形 模式 量化	量化 模式支持
服务器 CPU	fbgemm	支持的		所有 支持的
移动设备 CPU	qnnpack/xnnpack
服务器 GPU	TensorRT（早期 原型）	不支持 这个它 需要一个 图	支持的	静态 量化

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

• 支持 AVX2 或更高版本的 x86 CPU（不支持 AVX2 的话，某些操作的实现效率较低），通过 fbgemm

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

• (早期原型) 通过 TensorRT 实现对 NVidia GPU 的支持，通过 fx2trt (即将开源)

关于原生CPU后端的说明

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

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

fbgemm 的默认设置：

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

qnnpack 的默认设置：

# set the qconfig for PTQ
qconfig = torch.quantization.get_default_qconfig('qnnpack')
# or, set the qconfig for QAT
qconfig = torch.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	N
nn.MultiheadAttention	不支持	不支持
激活函数	广泛支持	未更改时， 计算 保持在 fp32

注意：这将很快用从本地 backend_config_dict 生成的一些信息进行更新。

量化 API 参考

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

量化后端配置

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

量化精度调试

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

量化定制

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

我们还为conv2d()、conv3d()和linear()提供按通道量化支持

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

量化自定义模块 API

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

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

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

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

4. 一个描述上述（1）、（2）、（3）的配置，传递给量化 API。

框架接下来将执行以下操作：

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

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

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

Example:

import torch
import torch.nn.quantized as nnq
import torch.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.quantization.default_qconfig
mp = torch.quantization.prepare(
    m, prepare_custom_config_dict=prepare_custom_config_dict)
# calibration (not shown)
mq = torch.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_dict = {'': torch.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.quantization.quantize_fx.prepare_fx(
    m, qconfig_dict, prepare_custom_config_dict=prepare_custom_config_dict)
# calibration (not shown)
mq = torch.quantization.quantize_fx.convert_fx(
    mp, convert_custom_config_dict=convert_custom_config_dict)

最佳实践

1. 如果你使用的是 fbgemm 后端，我们需要使用 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 for fbgemm 或 qnnpack 后端

常见错误

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

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

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

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

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = torch.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.quantization.DeQuantStub 对张量进行反量化。在急切模式量化中，这需要手动完成。一个端到端的示例：

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = torch.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.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 图模式量化用户指南来解决这个问题。