XNNPACK Delegate 內部

這是 ExecuTorch XNNPACK 後端 delegate 的高層次概覽。這個高效能 delegate 旨在降低 ExecuTorch 模型的 CPU 推理延遲。我們將簡要介紹 XNNPACK 庫，並探討 delegate 的整體架構和預期用例。

什麼是 XNNPACK？

XNNPACK 是一個高度最佳化的神經網路運算子庫，適用於 Android、iOS、Windows、Linux 和 macOS 環境中的 ARM、x86 和 WebAssembly 架構。這是一個開源專案，您可以在 github 上找到更多相關資訊。

什麼是 ExecuTorch delegate？

Delegate 是後端處理和執行 ExecuTorch 程式部分的入口點。ExecuTorch 模型中被 delegate 的部分會將執行權交給後端。XNNPACK 後端 delegate 是 ExecuTorch 中眾多可用的 delegate 之一。它利用 XNNPACK 第三方庫在各種 CPU 上高效加速 ExecuTorch 程式。關於 delegate 和開發自己的 delegate 的更多詳細資訊可在此處找到。此處。建議您在繼續閱讀架構部分之前熟悉這部分內容。

架構

提前 (Ahead-of-time)

在 ExecuTorch 匯出流程中，向 XNNPACK delegate 的 lowering 發生在 to_backend() 階段。在此階段，模型由 XnnpackPartitioner 進行分割槽。圖中被分割槽的部分被轉換為 XNNPACK 特定的圖表示，然後透過 flatbuffer 進行序列化。序列化的 flatbuffer 即可由 XNNPACK 後端在執行時反序列化並執行。

分割槽器 (Partitioner)

分割槽器由後端 delegate 實現，用於標記適合進行 lowering 的節點。XnnpackPartitioner 使用節點目標和模組元資料進行 lowering。可以在此處找到更多關於分割槽器的參考資料

基於模組的分割槽

source_fn_stack 嵌入在節點的元資料中，提供有關這些節點來源的資訊。例如，torch.nn.Linear 等模組在被捕獲並匯出 to_edge 時，會為其計算生成一組節點。與計算線性模組相關的節點組的 source_fn_stack 即為 torch.nn.Linear. 基於source_fn_stack` 的分割槽使我們能夠識別可透過 XNNPACK 進行 lowering 的節點組。

例如，捕獲 torch.nn.Linear 後，您會在與線性模組相關的 addmm 節點的元資料中找到以下鍵

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

基於運算子的分割槽

XnnpackPartitioner 也使用運算子目標進行分割槽。它遍歷圖並識別可降低到 XNNPACK 的單個節點。基於模組的分割槽的一個缺點是可能會跳過來自 decompositions 的運算子。例如，像 torch.nn.Hardsigmoid 這樣的運算子被分解為 add、muls、divs 和 clamps。雖然 hardsigmoid 本身不可降低，但我們可以降低分解後的運算子。依賴 source_fn_stack 元資料會跳過這些可降低的運算子，因為它們屬於一個不可降低的模組，因此為了提高模型效能，我們根據運算子目標以及 source_fn_stack 貪婪地降低運算子。

Passes

在任何序列化之前，我們會在子圖上應用 passes 來準備圖。這些 passes 本質上是圖轉換，有助於提高 delegate 的效能。下面我們將概述最主要的 passes 及其功能。所有 passes 的描述請參見此處

• Channels Last 重塑

  • ExecuTorch 張量在傳入 delegate 之前通常是連續的，而 XNNPACK 只接受 channels-last 記憶體佈局。此 pass 會最大限度地減少插入的 permutation 運算子數量，以便傳入 channels-last 記憶體格式。

• Conv1d 轉 Conv2d

  • 透過將 Conv1d 節點轉換為 Conv2d，使其能夠被 delegate 處理

• 卷積和 BN 融合

  • 將 batch norm 操作與前一個卷積節點融合

序列化

從模型中分割槽出可降低的子圖後，XNNPACK delegate 會預處理這些子圖，並透過 flatbuffer 為 XNNPACK 後端將其序列化。

序列化 Schema

XNNPACK delegate 使用 flatbuffer 進行序列化。為了提高執行時效能，XNNPACK delegate 的 flatbuffer schema 模仿了 XNNPACK 庫的圖級別 API 呼叫。序列化的資料是 XNNPACK API 的引數，這樣在執行時，可以透過連續呼叫 XNNPACK API 有效地建立 XNNPACK 執行圖。

執行時

XNNPACK 後端的執行時透過自定義的 init 和 execute 函式與 ExecuTorch 執行時互動。每個被 delegate 的子圖都包含在一個單獨序列化的 XNNPACK blob 中。模型初始化時，ExecuTorch 會對所有 XNNPACK Blobs 呼叫 init 函式，以從序列化的 flatbuffer 載入子圖。之後，執行模型時，每個子圖都通過後端透過自定義的 execute 函式執行。要詳細瞭解 delegate 執行時如何與 ExecuTorch 互動，請參閱此資源。

XNNPACK 庫

XNNPACK delegate 支援多個平臺上的 CPU；有關支援的硬體架構的更多資訊，請參閱 XNNPACK 庫的 README。

Init

呼叫 XNNPACK delegate 的 init 函式時，我們透過 flatbuffer 反序列化預處理過的 blob。我們使用提前序列化的資訊來定義節點（運算子）和邊（中間張量），以構建 XNNPACK 執行圖。正如我們之前提到的，大部分處理工作已提前完成，因此在執行時我們只需連續呼叫 XNNPACK API 並傳入序列化的引數。當我們將靜態資料定義到執行圖中時，XNNPACK 會在執行時執行權重打包，為權重和偏置等靜態資料的高效執行做準備。建立執行圖後，我們建立執行時物件並將其傳遞給 execute 函式。

由於權重打包會在 XNNPACK 內部建立權重的額外副本，我們會釋放預處理過的 XNNPACK Blob 中權重的原始副本，這有助於減少部分記憶體開銷。

Execute

執行 XNNPACK 子圖時，我們準備好張量輸入和輸出，並將它們饋送到 XNNPACK 執行時圖。執行執行時圖後，輸出指標將填充計算出的張量。

效能分析

我們已為 XNNPACK delegate 啟用了基本效能分析功能，可透過編譯器標誌 -DEXECUTORCH_ENABLE_EVENT_TRACER 啟用（新增 -DENABLE_XNNPACK_PROFILING 可獲得更多詳細資訊）。透過 ExecuTorch 的開發者工具整合，您現在也可以使用開發者工具對模型進行效能分析。您可以按照使用 ExecuTorch 開發者工具分析模型中的步驟，瞭解如何分析 ExecuTorch 模型以及如何使用開發者工具的 Inspector API 檢視 XNNPACK 的內部效能分析資訊。xnn_executor_runner 中提供了示例實現（請參閱此處教程）。

量化

XNNPACK delegate 也可用作執行對稱量化模型的後端。對於量化模型 delegate，我們使用 XNNPACKQuantizer 對模型進行量化。Quantizer 是後端特定的，這意味著 XNNPACKQuantizer 被配置為對模型進行量化，以利用 XNNPACK 庫提供的量化運算子。我們不會詳細介紹如何實現自定義 quantizer，您可以按照此處文件進行操作。但是，我們將簡要概述如何量化模型以利用 XNNPACK delegate 的量化執行。

配置 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 對模型進行量化

配置好我們的 quantizer 後，現在就可以對模型進行量化了

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 量化的更深入解釋。

另請參閱

• 在 Android AAR 中整合 XNNPACK Delegate

• 完成 Lowering 到 XNNPACK 教程