目录

量化

警告

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

量化简介

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

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

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

量化要求用户了解三个概念:

  1. 量化配置(Qconfig):指定权重和激活值的量化方式。需要 Qconfig 来创建量化模型。

  2. 后端:指支持量化操作的内核,通常具有不同的数值表示方式。

  3. 量化引擎(torch.backends.quantization.engine):当执行量化模型时,qengine 指定用于执行的后端。确保 qengine 与 Qconfig 一致非常重要。

原生支持的后端

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

相应的实现会根据PyTorch的构建模式自动选择,尽管用户可以通过将torch.backends.quantization.engine设置为fbgemmqnnpack来覆盖此设置。

注意

目前,PyTorch 还未在 CUDA 上提供量化操作符的实现 - 这是未来工作的方向。为了测试量化功能,请将模型移动到 CPU 上。

量化感知训练(通过 FakeQuantize, 该方法在fp32中模拟量化数值)支持CPU和CUDA。

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

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'

量化 API 概述

PyTorch 提供两种不同的量化模式:Eager Mode Quantization 和 FX Graph Mode Quantization。

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

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图模式量化,如果它不起作用,用户可以尝试遵循使用FX图模式量化的指南或回退到eager模式量化。

下表比较了 Eager Mode Quantization(急切模式量化)和 FX Graph Mode Quantization(FX 图模式量化)之间的差异:

Eager Mode 量化

FX 图形 模式 量化

发布 状态

测试版

原型

操作符 融合

手册

自动的

量化/反量化 放置

手册

自动的

量化模块

支持的

支持的

量化 函数/Torch 操作

手册

自动的

支持 自定义

有限支持

完全支持

量化模式 支持

训练后 量化: 静态、动态、 仅权重

量化感知训练: 静态

训练后 量化: 静态、动态、 仅权重

量化感知训练: 静态

输入/输出 模型类型

torch.nn.Module

torch.nn.Module (可能需要一些重构 以使模型与 FX 图模式 量化兼容)

目前支持三种量化类型:

  1. 动态量化(权重使用激活值以浮点数形式读取/存储,并对计算进行量化。)

  2. 静态量化(权重量化,激活量化,训练后需要校准)

  3. 静态量化感知训练(权重量化,激活值量化,在训练过程中建模量化数值)

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

操作符覆盖范围在动态量化和静态量化之间有所不同,如下表所示。 请注意,对于 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

急切模式量化

动态量化

这是最易于应用的量化形式,其中权重在事先进行量化,而激活值则在推理过程中动态量化。这种形式适用于模型执行时间主要由从内存中加载权重而非计算矩阵乘法所主导的情况。这种情况常见于具有小批量大小的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)

要了解更多关于动态量化的内容,请参阅我们的动态量化教程

静态量化

静态量化会对模型的权重和激活进行量化。它会在可能的情况下将激活操作融合到前一层中。它需要使用一个代表性的数据集进行校准,以确定激活的最佳量化参数。后训练量化通常在内存带宽和计算节省都重要的情况下使用,卷积神经网络(CNN)是典型的使用场景。静态量化也被称为后训练量化或PTQ。

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)

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

量化感知训练

量化感知训练在训练过程中模拟量化的影响,相比其他量化方法可以实现更高的精度。在训练期间,所有计算均以浮点数进行,通过 clamp 和四舍五入操作,fake_quant 模块模拟了 INT8 量化的效应。模型转换后,权重和激活值会被量化,并尽可能将激活值融合到前一层中。它通常与卷积神经网络(CNN)一起使用,相比静态量化能获得更高的精度。量化感知训练也被称为 QAT。

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教程

(原型) FX 图模式量化

FX图模式支持的量化类型可以分为两种方式:

  1. 训练后量化(在训练后应用量化,量化参数基于样本校准数据计算)

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

然后,这两者中的每一个可能包含以下任何一种或全部类型:

  • 仅权重量化(仅权重进行静态量化)

  • 动态量化(权重静态量化,激活值动态量化)

  • 静态量化(权重和激活值均被静态量化)

这两种分类方式是独立的,因此理论上我们可以有6种不同的量化类型。

FX图模式量化支持的量化类型有:

  • 训练后量化

    • 仅权重量化

    • 动态量化

    • 静态量化

  • 量化感知训练

    • 静态量化

在训练后量化中存在多种量化类型(仅权重、动态和静态),配置通过 qconfig_dictprepare_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)的更多信息:

量化张量

PyTorch 支持每张量和每通道的非对称线性量化。每张量表示张量内的所有值都以相同的方式进行缩放。每通道表示对于每个维度(通常是张量的通道维度),张量中的值会通过不同的值进行缩放和偏移(有效地使缩放和偏移变为向量)。这可以减少将张量转换为量化值时的误差。

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

_images/math-quantizer-equation.png

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

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

量化操作覆盖

量化张量仅支持常规全精度张量数据操作方法的一个有限子集。对于 PyTorch 中包含的神经网络运算符,我们限制支持为:

  1. 8 bit weights (data_type = qint8)

  2. 8 bit activations (data_type = quint8)

请注意,目前操作符实现仅支持对convlinear操作符的权重进行逐通道量化。此外,输入数据的最小值和最大值会线性映射到量化数据类型的最小值和最大值,以确保零值没有量化误差。

其他数据类型和量化方案可以通过自定义运算符机制实现。

许多针对量化张量的操作都与完整浮点版本在 torchtorch.nn 中的 API 一致。执行重新量化的神经网络模块的量化版本可在 torch.nn.quantized 中找到。这些操作在操作签名中明确接受输出量化参数(scale 和 zero_point)。

此外,我们还支持与常见融合模式相对应的融合版本,这些模式会影响量化: torch.nn.intrinsic.quantized

对于感知量化训练,我们支持在 torch.nn.qattorch.nn.intrinsic.qat 处准备了感知量化训练的模块

The list of supported operations is sufficient to cover typical CNN and RNN models

量化定制

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

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

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

量化前的模型准备

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

  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. 指定在哪里对激活值进行量化和反量化。这是通过 QuantStubDeQuantStub 模块完成的。

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

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

最佳实践

  1. 设置观察者的 reduce_range 参数为 True,如果你使用的是 fbgemm 后端。此参数通过将量化数据类型的范围减少 1 位,防止某些 int8 指令的溢出。

常见错误

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

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

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.quantization

此模块实现了您直接调用的函数,用于将模型从FP32转换为量化形式。例如, prepare() 用于训练后量化,以准备校准步骤, 并且 convert() 实际上将权重转换为int8,并将操作替换为其量化版本。还有其他辅助函数用于对输入模型进行量化以及执行诸如conv+relu之类的关键融合操作。

torch.nn.intrinsic

这个模块实现了组合(融合)的模块 conv + relu,之后可以对其进行量化。

torch.nn.intrinsic.qat

该模块实现了用于量化感知训练所需的融合操作版本。

torch.nn.intrinsic.quantized

该模块实现了融合操作(如 conv + relu)的量化版本。

torch.nn.qat

该模块实现了关键的nn模块的版本 Conv2d()Linear(),这些模块在FP32中运行,但应用了四舍五入以模拟 INT8量化的效果。

torch.nn.quantized

此模块实现了量化版本的nn层,例如 ~`torch.nn.Conv2d` 和 torch.nn.ReLU

torch.nn.quantized.dynamic

动态量化 Linear, LSTM, LSTMCell, GRUCell, 和 RNNCell

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源