• 教程 >
  • 如何为 PyTorch 2 导出量化编写 Quantizer
快捷方式

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

**作者**:Leslie FangWeiwen XiaJiong GongKimish PatelJerry Zhang

简介

(原型) PyTorch 2 导出训练后量化 介绍了 PyTorch 2 导出量化的整体 API,与 fx 图模式量化在 API 方面的主要区别在于,我们明确地指出量化针对的是特定的后端。因此,要使用新的流程,后端需要实现一个 Quantizer 类,该类编码:(1) 后端支持哪些量化运算符或模式 (2) 用户如何表达他们希望如何对浮点模型进行量化,例如,将整个模型量化为 int8 对称量化,或仅量化线性层等。

请参阅 此处了解新 API 和 Quantizer 的动机。

用于 XNNPACK 的现有量化器对象定义在 QNNPackQuantizer 中。

注解 API

Quantizer 使用注解 API 来传达不同算子/模式的量化意图。注解 API 主要由 QuantizationSpecQuantizationAnnotation 组成。

QuantizationSpec 用于传达张量如何量化的意图,例如数据类型、比特宽度、最小值、最大值、对称与非对称等。此外,QuantizationSpec 还允许量化器指定如何观察张量值,例如 MinMaxObserverHistogramObserver,或一些自定义观察器。

QuantizationAnnotationQuantizationSpec 对象组成,用于注解模式的输入张量和输出张量。注解输入张量等价于注解输入边,而注解输出张量等价于注解节点。QuantizationAnnotation 是一个带有几个字段的 dataclass

  • input_qspec_map 字段是 Dict 类,用于将每个输入张量(作为输入边)映射到一个 QuantizationSpec

  • output_qspec 字段表示用于注解输出张量的 QuantizationSpec

  • _annotated 字段指示该节点是否已被量化器注解。

总之,注解 API 要求量化器注解图的边(输入张量)或节点(输出张量)。现在,我们将逐步介绍如何在不同类型的 QuantizationSpec 中使用注解 API。

1. 注解常见算子模式

为了使用量化模式/算子,例如 quantized add,后端开发者将有意愿对模式的输入和输出进行量化(如 QuantizationSpec 所表达)。以下是如何使用注解 API 在量化工作流程中传达此意图的示例流程(以 add 算子为例)。

  • 步骤 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:为模式的输入和输出定义 QuantizationSpecQuantizationSpec 定义了 data typeqscheme 和其他关于用户如何观察或伪量化张量的意图的量化参数。

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 节点之后,在后续的量化流程中,HistogramObserver 将在准备阶段插入到其两个输入节点和一个输出节点。并且 HistogramObserver 将在转换阶段被 quantize 节点和 dequantize 节点替换。

2. 注解共享量化参数的算子

用户希望注解一个量化模型,其中量化参数可以在某些张量之间明确共享,这是很自然的。两个典型的用例是

  • 示例 1:一个例子是对于 add,让两个输入共享量化参数可以使算子的实现更容易。如果不使用 SharedQuantizationSpec,我们必须像上面第 1 节中的示例那样注解 add,其中 add 的两个输入具有不同的量化参数。

  • 示例 2:另一个例子是共享输入和输出之间的量化参数。这通常来自诸如 maxpoolaverage_poolconcat 等算子。

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

注意

  • 共享是可传递的

    由于以下原因,某些张量可能会有效地使用共享量化规范:

    • 两个节点/边被配置为使用 SharedQuantizationSpec

    • 存在一些节点的现有共享。

    例如,假设我们有两个 conv 节点 conv1conv2,并且它们都馈送到 cat 节点:cat([conv1_out, conv2_out], ...)。假设 conv1conv2 的输出和 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_outconv1_out 也将共享量化参数。在观察到的图中,您将看到以下内容

    conv1 -> obs -> cat
    conv2 -> obs   /
    

    并且两个 obs 将是相同的观察器实例。

  • 输入边是输入节点和使用该输入的节点之间的连接,因此它是一个 Tuple[Node, Node]

  • 输出值是 FX Node

现在,如果我们想使用 SharedQuantizationSpec 重写 add 注解示例以指示两个输入张量共享量化参数。我们可以将其 QuantizationAnnotation 定义为

  • 步骤 1:在 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 这样的算子,其输入和输出张量具有预定义和固定的比例/零点。FixedQParamsQuantizationSpec 专为这种用例而设计。要使用 FixedQParamsQuantizationSpec,用户需要显式地传入 scalezero_point 参数。

  • 步骤 1:在 FX 图中识别原始浮点模式。我们可以使用 QuantizationSpec 示例中介绍的相同方法来识别 sigmoid 模式。

  • 步骤 2:使用固定 scalezero_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函数,它接受ObserverOrFakeQuantizeObserverBaseFakeQuantizeBase)列表作为输入。从每个ObserverOrFakeQuantize对象中,用户可以获取scalezero point值。用户可以定义其关于如何基于观察者或伪量化实例计算的量化参数来推导出新的scalezero point值的启发式方法。

  • 步骤 3:定义DerivedQuantizationSpec对象,它接受以下输入:EdgeOrNode对象的列表。与每个EdgeOrNode对象对应的观察者将传递到derive_qparams_fn函数中;derive_qparams_fn函数;其他一些量化参数,例如dtypeqscheme

  • 步骤 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。为了更好地理解最终示例,以下是在示例中使用的类和实用程序函数

关于 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 实现发生更改,则表示形式可能会发生更改。但这可能是出乎意料的,因为建模用户通常假设当急切模式模型代码不发生变化时,他们应该在程序捕获后获得相同的模型表示形式。此问题的一个具体影响是,如果Quantizer基于识别aten IR 模式进行注释,则它可能在 PyTorch 版本更新后无法识别该模式,并且相同的急切模式浮点数可能会保持未量化。

建议:使用SubgraphMatcherWithNameNodeMap进行模式匹配

因此,我们建议用户通过SubgraphMatcherWithNameNodeMapSubgraphMatcher的改进版本,使查询用户想要注释的节点更容易)识别模式,通过捕获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 导出量化流程中。针对特定注释用例,给出了QuantizationSpecSharedQuantizationSpecFixedQParamsQuantizationSpecDerivedQuantizationSpec的示例。您可以使用XNNPACKQuantizer作为示例开始实现您自己的Quantizer。之后,请按照本教程实际量化您的模型。

文档

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

查看文档

教程

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

查看教程

资源

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

查看资源