如何为 PyTorch 2 Export 量化编写一个 Quantizer¶
创建日期：2023 年 7 月 28 日 | 最后更新日期：2024 年 8 月 1 日 | 最后验证日期：2024 年 11 月 5 日
作者：Leslie Fang, Weiwen Xia, Jiong Gong, Kimish Patel, Jerry Zhang
引言¶
(prototype) PyTorch 2 Export 训练后量化 介绍了 pytorch 2 export 量化的总体 API，与 fx 图模式量化的主要 API 区别在于，我们明确指出量化是针对特定后端的。因此，要使用新流程，后端需要实现一个 Quantizer 类，该类包含：(1). 后端支持的量化算子或模式有哪些 (2). 用户如何表达他们希望其浮点模型被量化的方式，例如，将整个模型量化为 int8 对称量化，或仅量化线性层等。
关于新 API 和 Quantizer 的动机，请参阅 此处。
为 XNNPACK 定义的现有量化器对象位于 QNNPackQuantizer 中。
注解 API¶
Quantizer 使用注解 API 来传达不同算子/模式的量化意图。注解 API 主要包括 QuantizationSpec 和 QuantizationAnnotation。
QuantizationSpec 用于传达张量将如何被量化的意图，例如 dtype、位宽、min、max 值、对称 vs. 非对称等。此外，QuantizationSpec 还允许量化器指定张量值应如何被观察，例如 MinMaxObserver、HistogramObserver 或一些自定义观察器。
QuantizationAnnotation 由 QuantizationSpec 对象组成，用于注解模式的输入张量和输出张量。注解输入张量等同于注解输入边，而注解输出张量等同于注解节点。QuantizationAnnotation 是一个 dataclass，包含几个字段：
input_qspec_map字段是Dict类型，用于将每个输入张量（作为输入边）映射到其QuantizationSpec。output_qspec字段表示用于注解输出张量的QuantizationSpec；_annotated字段指示此节点是否已被量化器注解。
总而言之，注解 API 要求量化器注解图的边（输入张量）或节点（输出张量）。现在，我们将提供一个分步教程，说明如何使用注解 API 和不同类型的 QuantizationSpec。
1. 注解常见算子模式¶
为了使用量化模式/算子，例如 quantized add，后端开发者会有意图量化（如 QuantizationSpec 所表达）模式的输入和输出。以下是量化工作流中使用注解 API 传达此意图的示例流程（以 add 算子为例）。
步骤 1：在 FX 图中识别原始浮点模式。识别此模式有几种方法：量化器可以使用模式匹配器来匹配算子模式；量化器可以从头到尾遍历节点并比较节点的 target 类型以匹配算子模式。在此示例中，我们可以使用 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 中为add节点的两个输入和一个输出创建的QuantizationSpec来创建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 节点后，在后续的量化流程中，将在 prepare 阶段在其两个输入节点和一个输出节点处插入 HistogramObserver。在 convert 阶段，HistogramObserver 将被替换为 quantize 节点和 dequantize 节点。
3. 注解固定量化参数的算子¶
注解量化模型的另一个典型用例是其量化参数预先已知的张量。例如，像 sigmoid 这样的算子，它在输入和输出张量处具有预定义和固定的 scale/zero_point 值。FixedQParamsQuantizationSpec 专为此用例设计。要使用 FixedQParamsQuantizationSpec，用户需要显式传入 scale 和 zero_point 参数。
步骤 1：在 FX 图中识别原始浮点模式。我们可以使用与
QuantizationSpec示例中介绍的相同方法来识别sigmoid模式。步骤 2：创建
FixedQParamsQuantizationSpec对象，其输入为固定的scale、zero_point值。这些值将用于在 convert 阶段创建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 并使用 Torchvision Resnet18 运行一个玩具示例。为了更好地理解最终示例，这里列出了示例中使用的类和实用函数：
QuantizationConfig 分别包含用于激活、权重和偏置的
QuantizationSpec。注解模型时，get_input_act_qspec、get_output_act_qspec、get_weight_qspec 和 get_bias_qspec 可用于从
QuantizationConfig获取特定模式的QuantizationSpec。
关于 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 实现发生变化，表示可能会改变。但这可能是意外的，因为模型用户通常假设当 eager mode 模型代码不变时，他们在程序捕获后也应该得到相同的模型表示。这个问题的具体影响是，如果一个 Quantizer 基于识别 aten IR 模式进行注解，那么在 PyTorch 版本更新后可能无法识别该模式，并且相同的 eager mode 浮点模型可能无法被量化。
建议：使用 SubgraphMatcherWithNameNodeMap 进行模式匹配¶
因此，我们建议人们通过 SubgraphMatcherWithNameNodeMap（SubgraphMatcher 的改进版本，使查询用户想要注解的节点更容易）来识别模式，通过捕获一个 torch IR 模式（使用与捕获浮点模型相同的程序捕获），而不是直接使用 aten IR 模式。
示例
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 算子 target 之外，我们没有一个好的方法来识别我们想要注解哪个消费者节点。
另一个需要注意的地方是，我们需要确保我们有一个详尽的示例列表（例如 2D、3D、4D 输入，真实输入 vs. 符号输入，training=True vs. training=False 等），以确保覆盖从 torch IR 模式捕获的所有可能的 aten IR 结果。
注意：未来我们可能会提供一些（模式，示例输入列表）或一些预生成的匹配器对象，以便人们可以直接使用它们。
结论¶
通过本教程，我们介绍了 PyTorch 2 中的新量化路径。用户可以学习如何使用 QuantizationAnnotation API 定义一个 BackendQuantizer 并将其集成到 PyTorch 2 Export 量化流程中。QuantizationSpec、SharedQuantizationSpec、FixedQParamsQuantizationSpec 和 DerivedQuantizationSpec 的示例针对特定的注解用例给出。您可以参考 XNNPACKQuantizer 作为示例来开始实现自己的 Quantizer。之后请按照本教程实际量化您的模型。
