如何为PyTorch 2导出量化编写``Quantizer``¶
Created On: Jul 28, 2023 | Last Updated: Aug 01, 2024 | Last Verified: Nov 05, 2024
作者：Leslie Fang, Weiwen Xia, Jiong Gong, Kimish Patel, Jerry Zhang
简介¶
`(原型) PyTorch 2导出的后训练量化 <https://pytorch.org/tutorials/prototype/pt2e_quant_ptq.html>`__介绍了PyTorch 2导出量化的整体API，与FX图模式量化在API上的主要区别是我们明确表示量化是针对特定后端的。因此，要使用新的流程，后端需要实现一个``Quantizer``类，该类编码：(1)后端支持的量化运算符或模式；(2)用户如何表达他们希望将浮点模型量化的方式，比如将整个模型量化为int8对称量化，或仅量化线性层等。
有关新API和``Quantizer``的动机，请参见`此处 <https://pytorch.org/tutorials/prototype/pt2e_quant_ptq.html#motivation-of-pytorch-2-export-quantization>`__。
为``XNNPACK``定义的现有量化器对象位于`QNNPackQuantizer <https://github.com/pytorch/pytorch/blob/main/torch/ao/quantization/pt2e/quantizer/xnnpack_quantizer.py>`__
注释API¶
``Quantizer``使用注释API来传递对不同运算符/模式的量化意图。注释API主要包括`QuantizationSpec <https://github.com/pytorch/pytorch/blob/1ca2e993af6fa6934fca35da6970308ce227ddc7/torch/ao/quantization/_pt2e/quantizer/quantizer.py#L38>`__和`QuantizationAnnotation <https://github.com/pytorch/pytorch/blob/07104ca99c9d297975270fb58fda786e60b49b38/torch/ao/quantization/_pt2e/quantizer/quantizer.py#L144>`__。
QuantizationSpec``用于传递张量将如何量化的意图，例如数据类型、位宽、最小值、最大值、对称量化或非对称量化等。此外，``QuantizationSpec``还允许量化器指定如何观察张量值，例如``MinMaxObserver、``HistogramObserver``或一些自定义观察器。
由``QuantizationSpec``对象组成的``QuantizationAnnotation``用于注释模式的输入张量和输出张量。注释输入张量相当于注释输入边缘，而注释输出张量相当于注释节点。QuantizationAnnotation``是一个具有几个字段的``dataclass：
input_qspec_map``字段是一个``Dict``类，用于将每个输入张量（作为输入边缘）映射到一个``QuantizationSpec。output_qspec``字段表示用于注释输出张量的``QuantizationSpec；``_annotated``字段指示此节点是否已被量化器注释。
总之，注释API要求量化器注释图的边缘（输入张量）或节点（输出张量）。接下来，我们将逐步介绍如何使用具有不同类型``QuantizationSpec``的注释API。
1. 注释常见的运算符模式¶
为了使用量化的模式/运算符，例如``quantized add``，后端开发人员将希望（通过``QuantizationSpec``表达的）量化模式的输入和输出。以下是一个示例流程（以``add``运算符为例），展示如何在量化工作流中通过注释API传达这一意图。
步骤1：在FX图中识别原始浮点模式。有几种识别此模式的方法：量化器可以使用模式匹配器来匹配运算符模式；量化器可以从头至尾遍历节点并比较节点的目标类型以匹配运算符模式。在该示例中，我们可以使用`get_source_partitions <https://github.com/pytorch/pytorch/blob/07104ca99c9d297975270fb58fda786e60b49b38/torch/fx/passes/utils/source_matcher_utils.py#L51>`__匹配此模式。原始浮点``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``定义了用户关于如何观察或假量化张量的意图，例如“数据类型”、“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``节点进行这样的注释后，在后续的量化流程中，``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 并运行一个 `玩具示例`_ 使用 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量化流程使用 pre 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模块或功能操作的实现发生变化，表示可能会改变。但这一点可能是意料之外的，因为建模用户通常假设当即时模式模型代码保持不变时，在程序捕获之后也应该获得相同的模型表示。此问题的具体影响在于，如果 Quantizer 基于识别 aten IR 模式进行注释，那么在PyTorch版本更新后可能无法识别模式，并且相同的即时模式浮点可能被留为未量化。
建议：使用 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。在此之后，请遵循 本教程 实际量化您的模型。
