ExecuTorch 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 程序。有关委托和开发您自己的委托的更多详细信息，请参阅此处。建议您在继续阅读“架构”部分之前先熟悉该内容。

架构

提前编译 (Ahead-of-time)

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

Partitioner

partitioner 由后端委托实现，用于标记适合降低的节点。XnnpackPartitioner 使用节点目标和模块元数据进行降低。有关 partitioner 的更多参考信息，请参阅此处

基于模块的分区

source_fn_stack 嵌入在节点的元数据中，并提供有关这些节点来源的信息。例如，当捕获和导出 to_edge 时，像 torch.nn.Linear 这样的模块会为其计算生成节点组。与计算线性模块关联的节点组随后具有 torch.nn.Linear 的 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

在任何序列化之前，我们在子图上应用 pass 以准备图。这些 pass 本质上是图转换，有助于提高委托的性能。我们在下面概述了最重要的 pass 及其功能。有关所有 pass 的描述，请参阅此处

• Channels Last Reshape

  • ExecuTorch 张量在传递给委托之前往往是连续的，而 XNNPACK 仅接受 channels-last 内存布局。此 pass 最小化了为传入 channels-last 内存格式而插入的排列运算符的数量。

• Conv1d to Conv2d

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

• Conv 和 BN 融合

  • 将批归一化操作与之前的卷积节点融合

序列化 (Serialiazation)

在从模型中分区出可降低的子图之后，XNNPACK 委托预处理这些子图，并通过 flatbuffer 为 XNNPACK 后端序列化它们。

序列化 Schema

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

运行时 (Runtime)

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

XNNPACK 库

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

Init

当调用 XNNPACK 委托的 init 时，我们通过 flatbuffer 反序列化预处理的 blob。我们定义节点（运算符）和边（中间张量），以使用我们提前序列化的信息构建 XNNPACK 执行图。正如我们前面提到的，大部分处理已提前完成，因此在运行时，我们只需连续调用带有序列化参数的 XNNPACK API 即可。当我们将静态数据定义到执行图中时，XNNPACK 在运行时执行权重打包，以准备静态数据（如权重和偏差）以进行高效执行。在创建执行图后，我们创建运行时对象并将其传递给 execute。

由于权重打包在 XNNPACK 内部创建了权重的额外副本，因此我们释放了预处理的 XNNPACK Blob 内部的原始权重副本，这使我们能够消除一些内存开销。

Execute

当执行 XNNPACK 子图时，我们准备张量输入和输出，并将它们馈送到 XNNPACK 运行时图。在执行运行时图后，输出指针将填充计算后的张量。

分析 (Profiling)

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

量化 (Quantization)

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)

您现在将看到模型的 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 应用程序

• 完成降低到 XNNPACK 教程