如何为 PyTorch 2 导出量化编写 Quantizer
创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日
作者：方磊，夏伟文，龚炯，帕特尔·基米什，张杰瑞
前提条件: ¶
需要：
可选：
简介
（原型）PyTorch 2 导出训练后量化引入了 pytorch 2 导出量化的整体 API，与 FX 图模式量化在 API 方面的主要区别在于，我们明确指出量化是针对特定后端的。因此，要使用新的流程，后端需要实现一个 Quantizer 类，该类编码以下内容：（1）后端支持哪些量化算子或模式；（2）用户如何表达他们希望浮点模型进行量化的方式，例如，将整个模型量化为 int8 对称量化，或者仅量化线性层等。
请参阅此处了解新 API 的动机以及 Quantizer 。
已为 XNNPACK 定义的现有量化对象位于 QNNPackQuantizer 中
注解 API
Quantizer 使用注解 API 来传达不同算子/模式的量化意图。注解 API 主要由 QuantizationSpec 和 QuantizationAnnotation 组成。
QuantizationSpec 用于传达张量量化意图，例如数据类型、位宽、最小值、最大值、对称与非对称等。此外， QuantizationSpec 还允许量化器指定张量值应该如何观察，例如 MinMaxObserver ，或 HistogramObserver ，或一些自定义观察者。
QuantizationAnnotation 由 QuantizationSpec 对象组成，用于注释模式输入张量和输出张量。注释输入张量相当于注释输入边，而注释输出张量相当于注释节点。 QuantizationAnnotation 是一个具有多个字段的 dataclass 。
input_qspec_map字段是类Dict，用于将每个输入张量（作为输入边）映射到QuantizationSpec。output_qspec字段表示用于注释输出张量的QuantizationSpec。_annotated字段表示此节点是否已被量化器标注。
总结来说，标注 API 需要量化器标注图中的边（输入张量）或节点（输出张量）。现在，我们将逐步介绍如何使用标注 API 与不同类型的 QuantizationSpec 进行操作。
1. 注解常见操作模式
为了使用量化模式/操作符，例如 quantized add ，后端开发者会有量化（如通过 QuantizationSpec 表达）模式输入、输出的意图。以下是一个示例流程（以 add 操作符为例），说明这种意图如何在标注 API 的量化工作流程中传达。
步骤 1：在 FX 图中识别原始浮点模式。有几种方法可以识别这种模式：量化器可能使用模式匹配器来匹配操作符模式；量化器可能从开始到结束遍历节点，并将节点的目标类型与操作符模式进行比较。在这个例子中，我们可以使用 get_source_partitions 来匹配这种模式。原始浮点
add模式仅包含一个add节点。
add_partitions = get_source_partitions(gm.graph, [operator.add, torch.add])
add_partitions = list(itertools.chain(*add_partitions.values()))
for add_partition in add_partitions:
add_node = add_partition.output_nodes[0]
步骤 2：定义模式的输入和输出
QuantizationSpec。QuantizationSpec定义了data type、qscheme以及其他关于用户如何观察或模拟量化张量的量化参数。
act_quantization_spec = QuantizationSpec(
dtype=torch.int8,
quant_min=-128,
quant_max=127,
qscheme=torch.per_tensor_affine,
is_dynamic=False,
observer_or_fake_quant_ctr=HistogramObserver.with_args(eps=2**-12),
)
input_act_qspec = act_quantization_spec
output_act_qspec = act_quantization_spec
步骤 3：用
QuantizationAnnotation注释模式的输入和输出。在这个例子中，我们将使用步骤 2 中创建的QuantizationSpec对象为add节点的两个输入和一个输出创建QuantizationAnnotation对象。
input_qspec_map = {}
input_act0 = add_node.args[0]
input_qspec_map[input_act0] = input_act_qspec
input_act1 = add_node.args[1]
input_qspec_map[input_act1] = input_act_qspec
add_node.meta["quantization_annotation"] = QuantizationAnnotation(
input_qspec_map=input_qspec_map,
output_qspec=output_act_qspec,
_annotated=True,
)
在我们这样注释 add 节点之后，在随后的上量化流程中， HistogramObserver 将在准备阶段插入其两个输入节点和一个输出节点。在转换阶段， HistogramObserver 将被 quantize 节点和 dequantize 节点替换。
3. 注释具有固定量化参数的算子 ¶
另一个典型的注释量化模型的用例是对于量化参数事先已知的张量。例如，像 sigmoid 这样的算子，在输入和输出张量上具有预定义和固定的缩放/零点。FixedQParamsQuantizationSpec 是为这种用例设计的。要使用 FixedQParamsQuantizationSpec ，用户需要显式传递 scale 和 zero_point 的参数。
步骤 1：在 FX 图中识别原始浮点模式。我们可以使用在
QuantizationSpec示例中介绍的方法来识别sigmoid模式。步骤 2：创建具有固定
scale、zero_point值输入的FixedQParamsQuantizationSpec对象。这些值将用于在转换阶段创建quantize节点和dequantize节点。步骤 3：标注输入和输出以使用此
FixedQParamsQuantizationSpec对象。
act_qspec = FixedQParamsQuantizationSpec(
dtype=torch.uint8,
quant_min=0,
quant_max=255,
qscheme=torch.per_tensor_affine,
scale=1.0 / 256.0,
zero_point=0,
)
sigmoid_node.meta["quantization_annotation"] = QuantizationAnnotation(
input_qspec_map={input_act: act_qspec},
output_qspec=act_qspec,
_annotated=True,
)
4. 注释具有派生量化参数的张量 ¶
另一个用例是为从其他张量派生出的量化参数的张量定义约束。例如，如果我们想注释一个卷积节点，并定义其偏置输入张量的 scale 为激活张量的 scale 和权重张量的 scale 的乘积。我们可以使用 DerivedQuantizationSpec 来注释这个卷积节点。
步骤 1：在 FX 图中识别原始浮点模式。我们可以使用在
QuantizationSpec示例中介绍的方法来识别convolution模式。步骤 2：定义
derive_qparams_fn函数，它接受列表形式的ObserverOrFakeQuantize（ObserverBase 或 FakeQuantizeBase）作为输入。从每个ObserverOrFakeQuantize对象中，用户可以获取scale、zero point值。用户可以定义自己的启发式方法，根据从观察器或模拟量化实例计算出的量化参数推导出新的scale、zero point值。步骤 3：定义
DerivedQuantizationSpec对象，它接受输入为：列表形式的EdgeOrNode对象。与每个EdgeOrNode对象对应的观察器将被传递到derive_qparams_fn函数；derive_qparams_fn函数；以及其他几个量化参数，如dtype、qscheme。第 4 步：使用
QuantizationAnnotation注释此卷积节点的输入和输出。
def derive_qparams_fn(obs_or_fqs: List[ObserverOrFakeQuantize]) -> Tuple[Tensor, Tensor]:
assert len(obs_or_fqs) == 2, \
"Expecting two obs/fqs, one for activation and one for weight, got: {}".format(len(obs_or_fq))
act_obs_or_fq = obs_or_fqs[0]
weight_obs_or_fq = obs_or_fqs[1]
act_scale, act_zp = act_obs_or_fq.calculate_qparams()
weight_scale, weight_zp = weight_obs_or_fq.calculate_qparams()
return torch.tensor([act_scale * weight_scale]).to(torch.float32), torch.tensor([0]).to(torch.int32)
bias_qspec = DerivedQuantizationSpec(
derived_from=[(input_act, node), (weight, node)],
derive_qparams_fn=derive_qparams_fn,
dtype=torch.int32,
quant_min=-2**31,
quant_max=2**31 - 1,
qscheme=torch.per_tensor_symmetric,
)
input_qspec_map = {input_act: act_quantization_spec, weight: weight_quantization_spec, bias: bias_qspec}
node.meta["quantization_annotation"] = QuantizationAnnotation(
input_qspec_map=input_qspec_map,
output_qspec=act_quantization_spec,
_annotated=True,
)
5. Resnet18 的玩具示例 ¶
在定义了 QuantizationAnnotation API 注释方法之后，我们现在可以将它们组合起来构建 BackendQuantizer 并运行一个玩具示例。为了更好地理解最终示例，以下是示例中使用的类和实用函数：
QuantizationConfig 由用于激活、权重和偏置的
QuantizationSpec组成。在标注模型时，可以使用 get_input_act_qspec、get_output_act_qspec、get_weight_qspec 和 get_bias_qspec 来获取特定模式的
QuantizationSpec从QuantizationConfig。
关于 PT2E 量化流程的 IR（索引表示）笔记
IR 代表模型的中间表示，例如， torch IR（ torch.nn 模块， torch.nn.functional 操作）或 aten IR（ torch.ops.aten.linear ，…）。PT2E 量化流程使用预 autograd aten IR（torch.export API 的输出）以支持训练。正如之前所展示的，在它们上附加注释之前，我们需要匹配操作或操作模式，所以问题是我们要如何匹配模式？
动机：直接匹配 aten IR 的问题
最直接的方法可能是直接匹配 aten IR。
例如：
for n in gm.graph.nodes:
if n.op != "call_function" or n.target not in [
torch.ops.aten.relu.default,
torch.ops.aten.relu_.default,
]:
continue
relu_node = n
maybe_conv_node = n.args[0]
if (
not isinstance(maybe_conv_node, Node)
or maybe_conv_node.op != "call_function"
or maybe_conv_node.target
not in [
torch.ops.aten.conv1d.default,
torch.ops.aten.conv2d.default,
]
):
continue
# annotate conv and relu nodes
...
然而，使用此 IR 的问题之一是，如果 PyTorch 实现模块或功能操作发生变化，其表示可能也会改变。但这种情况可能是意外的，因为建模用户通常假设在急切模式模型代码没有变化的情况下，他们在程序捕获后应该得到相同的模型表示。这个问题的一个具体影响是，如果基于识别 aten IR 模式进行 Quantizer 注释，那么在 PyTorch 版本更新后可能无法识别该模式，并且相同的急切模式浮点数可能未被量化。
建议：使用 SubgraphMatcherWithNameNodeMap 进行模式匹配
因此，我们建议人们通过 SubgraphMatcherWithNameNodeMap （ SubgraphMatcher 的改进版本，使得查询人们想要标注的节点更加容易），通过捕获一个 torch 索引模式（使用与捕获浮点模型相同的程序捕获），而不是直接使用 aten 索引模式。
例如：
def conv_relu_pattern(input, weight, bias):
conv = torch.nn.functional.conv2d(input, weight, bias)
output = torch.nn.functional.relu(conv)
# returns an additional dict that includes a map from name to node that we want to annotate
return relu, {"input": input, "weight": weight, "bias": bias, "output": output}
matcher = SubgraphMatcherWithNameNodeMap(conv_relu_pattern)
matches = matcher.match(model)
for match in matches:
# find input and output of the pattern
# annotate the nodes
name_node_map = match.name_node_map
input_node = name_node_map["input"]
weight_node = name_node_map["weight"]
bias_node = name_node_map["bias"]
output_node = name_node_map["relu"]
input_node.users[0].meta["quantization_annotation"] = ...
weight_node.users[0].meta["quantization_annotation"] = ...
bias_node.users[0].meta["quantization_annotation"] = ...
output_node.meta["quantization_annotation"] = ...
这样，即使 nn 模块和函数的实现发生变化， Quantizer 仍然有效， aten 浮点模型的 IR 会发生变化，但由于我们重新捕获了模式而不是硬编码模式的 aten IR，我们将获得更新的 aten IR，并且仍然能够匹配该模式。
一个需要注意的问题是，如果模式的输入有多个用户，我们除了检查 aten 操作的目标之外，没有很好的方法来识别我们想要标注的用户节点。
另一个注意事项是，我们需要确保有一个详尽的示例列表（例如，2D、3D、4D 输入，实数与符号输入，训练=True 与训练=False 等），以确保模式能够覆盖从 torch IR 模式捕获的不同可能的 aten IR 结果。
注意：我们可能会提供一些（模式，示例输入列表）或一些预生成的匹配器对象，以便人们可以直接在未来使用它们。
结论 ¶
本教程介绍了 PyTorch 2 中的新量化路径。用户可以了解如何使用 QuantizationAnnotation API 定义 BackendQuantizer 并将其集成到 PyTorch 2 导出量化流程中。针对特定的注释用例给出了 QuantizationSpec 、 SharedQuantizationSpec 、 FixedQParamsQuantizationSpec 和 DerivedQuantizationSpec 的示例。您可以使用 XNNPACKQuantizer 作为示例来开始实现您自己的 Quantizer 。之后请按照本教程进行操作，以实际量化您的模型。
