ExecuTorch XNNPACK 委托¶

这是 ExecuTorch XNNPACK 后端委托的高级概述。该高性能委托旨在降低 ExecuTorch 模型在 CPU 上的推理延迟。我们将简要介绍 XNNPACK 库，并探讨该委托的整体架构及其预期用例。

什么是 XNNPACK？¶

XNNPACK 是一个高度优化的神经网络算子库，适用于 ARM、x86 和 WebAssembly 架构，支持 Android、iOS、Windows、Linux 和 macOS 环境。它是一个开源项目，您可以在 GitHub 上找到更多信息。

什么是 ExecuTorch 委托？¶

委托是后端处理和执行ExecuTorch程序部分的入口点。ExecuTorch模型的委托部分将执行交给后端。XNNPACK后端委托是ExecuTorch中可用的众多委托之一。它利用XNNPACK第三方库在各种CPU上高效加速ExecuTorch程序。关于委托和开发您自己的委托的更详细信息，请参阅这里。建议您在继续阅读架构部分之前熟悉这些内容。

架构¶

High Level XNNPACK delegate Architecture

Ahead-of-time¶

在 ExecuTorch 导出流程中，降低到 XNNPACK 委托发生在 to_backend() 阶段。在此阶段，模型由 XnnpackPartitioner 进行分区。图的分区部分被转换为 XNNPACK 特定的图表示形式，然后通过 flatbuffer 进行序列化。序列化的 flatbuffer 随后准备好在运行时由 XNNPACK 后端反序列化并执行。

ExecuTorch XNNPACK delegate Export Flow

分区器¶

分区器由后端委托实现，用于标记适合降低的节点。 XnnpackPartitioner 使用节点目标和模块元数据进行降低。更多关于分区器的参考资料可以在这里找到。

基于模块的分区¶

source_fn_stack 嵌入在节点的元数据中，并提供有关这些节点来源的信息。例如，像 torch.nn.Linear 这样的模块在被捕获并导出时 to_edge 会为其计算生成节点组。与计算线性模块相关的节点组随后具有 source_fn_stack torch.nn.Linear. Partitioning based on source_fn_stack` 使我们能够识别可通过 XNNPACK 降低的节点组。

例如，在捕获torch.nn.Linear之后，您会在与 linear 关联的 addmm 节点的元数据中找到以下键：

>>> print(linear_node.meta["source_fn_stack"])
'source_fn_stack': ('fn', <class 'torch.nn.modules.linear.Linear'>)

基于操作的分区¶

The XnnpackPartitioner 也使用操作目标进行分区。它遍历图并识别可以降低到XNNPACK的单个节点。基于模块的分区的一个缺点是，来自分解的操作可能会被跳过。例如，像torch.nn.Hardsigmoid这样的操作会被分解为加法、乘法、除法和夹紧操作。虽然hardsigmoid不能被降低，但我们可以降低分解后的操作。依赖于source_fn_stack元数据会跳过这些可降低的操作，因为它们属于不可降低的模块，因此为了提高模型性能，我们根据操作目标以及source_fn_stack贪婪地降低操作。

传递¶

在进行任何序列化之前，我们会对子图应用一些传递以准备图。这些传递本质上是图变换，有助于提高委托的性能。我们在下面概述了最重要的传递及其功能。有关所有传递的描述，请参见这里：

通道优先重塑
- 在执行 Torch 张量传入委托之前通常要求其为连续内存布局，而 XNNPACK 仅接受通道优先（channels-last）的内存布局。此转换步骤旨在最小化为适配通道优先内存格式所插入的置换算子数量。
Conv1d 转 Conv2d
- 允许我们通过将 Conv1d 节点转换为 Conv2d 来委托处理。
卷积与批归一化融合
- 将批归一化操作与前面的卷积节点融合

序列化¶

在从模型中分割出可下放的子图后，XNNPACK 委托会对这些子图进行预处理，并通过 flatbuffer 将其序列化以供 XNNPACK 后端使用。

序列化模式¶

XNNPACK委托使用flatbuffer进行序列化。为了提高运行时性能，XNNPACK委托的flatbuffer 模式镜像了XNNPACK库的图形级别API调用。序列化的数据是XNNPACK API的参数，因此在运行时，可以通过连续调用XNNPACK的API高效地创建XNNPACK执行图。

运行时¶

XNNPACK 后端的运行时通过自定义 init 和 execute 函数与 ExecuTorch 运行时进行接口交互。每个委托的子图都包含在一个单独序列化的 XNNPACK blob 中。当模型初始化时，ExecuTorch 会对所有 XNNPACK Blobs 调用 init 以从序列化的 flatbuffer 加载子图。之后，在模型执行时，每个子图通过后端和自定义 execute 函数执行。要了解更多关于委托运行时如何与 ExecuTorch 接口的信息，请参阅此资源。

XNNPACK 库¶

XNNPACK 委托支持多个平台的 CPU；有关支持的硬件架构的更多信息，请参阅 XNNPACK 库的 README。

初始化¶

在调用 XNNPACK 委托的init时，我们通过 flatbuffer 反序列化预处理的 blob。我们利用预先序列化的信息来定义节点（算子）和边（中间张量），以构建 XNNPACK 执行图。正如前面提到的，大部分处理工作已提前完成，因此在运行时，我们可以依次使用序列化的参数调用 XNNPACK API。由于我们将静态数据定义到了执行图中，XNNPACK 会在运行时进行权重打包，以准备权重和偏置等静态数据，从而实现高效执行。创建执行图后，我们创建运行时对象并将其传递给execute。

由于权重打包会在 XNNPACK 内部创建一份额外的权重副本，我们释放预处理后的 XNNPACK Blob 中原有的权重副本，从而降低部分内存开销。

执行¶

在执行 XNNPACK 子图时，我们准备张量输入和输出，并将其提供给 XNNPACK 运行时图。执行运行时图后，输出指针将被填充为计算得到的张量。

调试¶

我们已启用基本的XNNPACK委托分析功能，可以通过以下编译器标志 -DENABLE_XNNPACK_PROFILING 启用。通过ExecuTorch的开发者工具集成，您现在也可以使用开发者工具来分析模型。您可以按照使用ExecuTorch开发者工具分析模型中的步骤来分析ExecuTorch模型，并使用开发者工具的Inspector API查看XNNPACK的内部分析信息。

量化¶

XNNPACK 委托还可以用作后端来执行对称量化模型。对于量化模型委托，我们使用 XNNPACKQuantizer 量化模型。Quantizers 是特定于后端的，这意味着 XNNPACKQuantizer 被配置为量化模型以利用 XNNPACK 库提供的量化操作符。我们不会详细介绍如何实现自定义量化器，您可以按照文档这里进行操作。然而，我们将提供一个简要概述，说明如何量化模型以利用 XNNPACK 委托的量化执行。

配置 XNNPACKQuantizer¶

from torch.ao.quantization.quantizer.xnnpack_quantizer import (
  XNNPACKQuantizer,
  get_symmetric_quantization_config,
)
quantizer = XNNPACKQuantizer()
quantizer.set_global(get_symmetric_quantization_config())

在此处初始化XNNPACKQuantizer，并将量化配置设置为对称量化。对称量化是指权重使用qmin = -127和qmax = 127进行对称量化，这将强制量化零点为零。get_symmetric_quantization_config()可以使用以下参数进行配置：

is_per_channel
- 权重按通道进行量化
is_qat
- 感知量化的训练
is_dynamic
- 动态量化

然后我们可以根据需要配置XNNPACKQuantizer。下面我们将以下配置作为示例进行设置：

quantizer.set_global(quantization_config)
    .set_object_type(torch.nn.Conv2d, quantization_config) # can configure by module type
    .set_object_type(torch.nn.functional.linear, quantization_config) # or torch functional op typea
    .set_module_name("foo.bar", quantization_config)  # or by module fully qualified name

使用 XNNPACKQuantizer 量化您的模型¶

配置好量化器后，我们现在就可以对模型进行量化了。

from torch.export import export_for_training

exported_model = export_for_training(model_to_quantize, example_inputs).module()
prepared_model = prepare_pt2e(exported_model, quantizer)
print(prepared_model.graph)

Prepare 执行一些 Conv2d-BN 融合操作，并在适当位置插入量化观察者。对于训练后量化，我们通常在此步骤之后对模型进行校准。我们通过 prepared_model 运行示例样本，以观察张量的统计信息并计算量化参数。

最后，我们在此转换我们的模型：

quantized_model = convert_pt2e(prepared_model)
print(quantized_model)

您现在将看到模型的量化/反量化表示，这意味着在量化算子输入处插入了torch.ops.quantized_decomposed.dequantize_per_tensor，在算子输出处插入了torch.ops.quantized_decomposed.quantize_per_tensor。示例：

def _qdq_quantized_linear(
    x_i8, x_scale, x_zero_point, x_quant_min, x_quant_max,
    weight_i8, weight_scale, weight_zero_point, weight_quant_min, weight_quant_max,
    bias_fp32,
    out_scale, out_zero_point, out_quant_min, out_quant_max
):
    x_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor(
        x_i8, x_scale, x_zero_point, x_quant_min, x_quant_max, torch.int8)
    weight_fp32 = torch.ops.quantized_decomposed.dequantize_per_tensor(
        weight_i8, weight_scale, weight_zero_point, weight_quant_min, weight_quant_max, torch.int8)
    out_fp32 = torch.ops.aten.linear.default(x_fp32, weight_fp32, bias_fp32)
    out_i8 = torch.ops.quantized_decomposed.quantize_per_tensor(
        out_fp32, out_scale, out_zero_point, out_quant_min, out_quant_max, torch.int8)
    return out_i8

您可以阅读更多关于PyTorch 2量化深入解释这里。