XNNPACK 委托内部机制

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

什么是 XNNPACK？

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

什么是 ExecuTorch 委托？

委托是后端处理和执行 ExecuTorch 程序一部分的入口点。ExecuTorch 模型的委托部分会将执行交给后端。XNNPACK 后端委托是 ExecuTorch 中提供的众多委托之一。它利用 XNNPACK 第三方库来高效地加速各种 CPU 上的 ExecuTorch 程序。有关委托和开发您自己的委托的更详细信息，请在此 处 查看。建议在继续学习架构部分之前，先熟悉该内容。

架构

提前编译

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

分区器

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

基于模块的分区

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

例如，在捕获 torch.nn.Linear 后，您将在与线性相关的 addmm 节点的元数据中找到以下键：

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

基于操作的分区

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

传递

在进行任何序列化之前，我们会在子图上应用传递（passes）来准备图。这些传递本质上是图转换，有助于提高委托的性能。我们将在下方概述最显著的传递及其功能。有关所有传递的描述，请参见 此处。

• 通道优先重塑

  • ExecuTorch 张量在传递给委托之前通常是连续的，而 XNNPACK 只接受通道优先（channels-last）内存布局。此传递最大限度地减少了为通道优先内存格式插入的排列算子的数量。

• Conv1d 转 Conv2d

  • 通过将 Conv1d 节点转换为 Conv2d，允许我们委托 Conv1d 节点。

• Conv 和 BN 融合

  • 将批归一化（batch norm）操作与之前的卷积节点融合。

序列化

在从模型中分区出可降低的子图后，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 委托启用了基本的性能分析，可以通过编译器标志 -DEXECUTORCH_ENABLE_EVENT_TRACER (添加 -DENABLE_XNNPACK_PROFILING 获取更多详细信息) 来启用。通过 ExecuTorch 的开发者工具集成，您现在还可以使用开发者工具来分析模型。您可以遵循 使用 ExecuTorch 开发者工具分析模型 中的步骤来分析 ExecuTorch 模型，并使用开发者工具的 Inspector API 查看 XNNPACK 的内部性能分析信息。在 executor_runner 中提供了一个示例实现（请参阅 此处的教程）。

量化

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

配置 XNNPACKQuantizer

from executorch.backends.xnnpack.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 融合，并在适当的位置插入量化观察器。对于训练后量化（Post-Training Quantization），我们通常在此步骤之后校准模型。我们将样本示例通过 prepared_model 运行，以观察张量的统计数据来计算量化参数。

最后，我们在此处转换模型。

quantized_model = convert_pt2e(prepared_model)
print(quantized_model)

您现在将看到模型的 Q/DQ 表示，这意味着 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 量化的更深入的解释：这里。

另请参阅

• 将 XNNPACK 委托集成到 Android AAR

• 完成 XNNPACK 委托的降低教程