如何为PyTorch 2导出量化编写`Quantizer`¶

创建于：2023年7月28日 | 最后更新：2024年8月1日 | 最后验证：2024年11月5日

作者: Leslie Fang, Weiwen Xia, Jiong Gong, Kimish Patel, Jerry Zhang

先决条件:¶

必填项：

可选的：

介绍¶

(prototype) PyTorch 2 Export Post Training Quantization 介绍了 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传达此意图。

第一步：在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]

第二步：为模式的输入和输出定义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节点。

2. 注释共享量化参数的运算符¶

用户自然希望注释一个量化模型，其中量化参数可以在一些张量之间显式共享。两个典型的用例是：

示例1：一个例子是add，其中两个输入共享量化参数使得操作符的实现更加容易。如果不使用SharedQuantizationSpec，我们必须像上面第1节中的示例那样注释add，其中add的两个输入具有不同的量化参数。
示例2：另一个例子是在输入和输出之间共享量化参数。这通常是由诸如maxpool、average_pool、concat等运算符引起的。

SharedQuantizationSpec 是为这种情况设计的，用于注释其量化参数与其他张量共享的张量。SharedQuantizationSpec 的输入是一个 EdgeOrNode 对象，它可以是输入边或输出值。

注意

共享是可传递的

一些张量可能由于以下原因有效地使用了共享量化规范：
- 两个节点/边被配置为使用SharedQuantizationSpec。
- 存在某些节点的现有共享。
例如，假设我们有两个conv节点conv1和conv2，它们都被输入到一个cat节点：cat([conv1_out, conv2_out], ...)。假设conv1、conv2的输出以及cat的第一个输入都配置了相同的QuantizationSpec配置。cat的第二个输入配置为使用SharedQuantizationSpec与第一个输入共享。
```
conv1_out: qspec1(dtype=torch.int8, ...)
conv2_out: qspec1(dtype=torch.int8, ...)
cat_input0: qspec1(dtype=torch.int8, ...)
cat_input1: SharedQuantizationSpec((conv1, cat))  # conv1 node is the first input of cat
```
首先，conv1的输出隐式地与cat的第一个输入共享量化参数（和观察者对象），同样的情况也适用于conv2的输出和cat的第二个输入。因此，由于用户配置了cat的两个输入以共享量化参数，通过传递性，conv2_out和conv1_out也将共享量化参数。在观察到的图中，您将看到以下内容：
```
conv1 -> obs -> cat
conv2 -> obs   /
```
并且两个 obs 将是相同的观察者实例。

输入边是输入节点和消费输入的节点之间的连接，因此它是一个Tuple[Node, Node]。
输出值是一个FX Node。

现在，如果我们想用SharedQuantizationSpec重写add注释示例，以指示两个输入张量共享量化参数。我们可以这样定义它的QuantizationAnnotation：

第一步：在FX图中识别原始的浮点模式。我们可以使用与QuantizationSpec示例中介绍的相同方法来识别add模式。
步骤2：使用QuantizationSpec注释add的input_act0。
步骤3：创建一个SharedQuantizationSpec对象，输入边定义为(input_act0, add_node)，这意味着共享用于此边的观察者。然后，用户可以使用这个SharedQuantizationSpec对象来注释input_act1。

input_qspec_map = {}
share_qparams_with_input_act0_qspec = SharedQuantizationSpec((input_act0, add_node))
input_qspec_map = {input_act0: act_quantization_spec, input_act1: share_qparams_with_input_act0_qspec}

add_node.meta["quantization_annotation"] = QuantizationAnnotation(
    input_qspec_map=input_qspec_map,
    output_qspec=act_quantization_spec,
    _annotated=True,
)

3. 使用固定量化参数注释运算符¶

另一个典型的用例是为那些量化参数事先已知的张量进行注释。例如，像sigmoid这样的操作符，它在输入和输出张量上具有预定义和固定的scale/zero_point。 FixedQParamsQuantizationSpec 就是为这种用例设计的。要使用FixedQParamsQuantizationSpec，用户需要显式传入scale和zero_point的参数。

步骤1：在FX图中识别原始的浮点模式。我们可以使用与QuantizationSpec示例中介绍的相同方法来识别sigmoid模式。
第二步：使用固定的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 来注释这个卷积节点。

第一步：在FX图中识别原始的浮点模式。我们可以使用QuantizationSpec示例中介绍的相同方法来识别convolution模式。
第二步：定义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运行一个toy example。为了更好地理解最终的示例，以下是示例中使用的类和实用函数：

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 量化流程使用预自动微分 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模式模型代码没有变化时，他们应该在程序捕获后获得相同的模型表示。这个问题的一个具体影响是，如果Quantizer基于识别aten IR模式进行注释，那么在PyTorch版本更新后，它可能无法识别该模式，并且相同的eager模式浮点数可能未被量化。

建议：使用`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操作目标外，我们没有好的方法来识别我们想要注释的用户节点。

另一个注意事项是，我们需要确保我们有一个详尽的示例列表（例如2D、3D、4D输入，真实与符号输入，training=True与training=False等），以确保覆盖从torch IR模式捕获的不同可能的aten IR结果。

注意：我们可能会提供一些（模式，示例输入列表）或一些预生成的匹配器对象，以便人们将来可以直接使用它们。

结论¶

通过本教程，我们介绍了PyTorch 2中的新量化路径。用户可以学习如何使用QuantizationAnnotation API定义BackendQuantizer并将其集成到PyTorch 2导出量化流程中。针对特定的注释用例，提供了QuantizationSpec、SharedQuantizationSpec、FixedQParamsQuantizationSpec和DerivedQuantizationSpec的示例。您可以使用XNNPACKQuantizer作为示例开始实现自己的Quantizer。之后，请按照本教程实际量化您的模型。

如何为PyTorch 2导出量化编写Quantizer¶