模型匯出和降低

本節介紹將 PyTorch 模型轉換為 ExecuTorch 使用的執行時格式的過程。此過程通常稱為“匯出”，因為它使用 PyTorch 的匯出功能將 PyTorch 模型轉換為適合在裝置上執行的格式。此過程會生成一個 .pte 檔案，該檔案已針對特定後端進行了裝置上執行最佳化。

先決條件

匯出需要安裝 ExecuTorch Python 庫，通常透過執行 pip install executorch 來完成。有關更多資訊，請參閱 安裝。此過程假設您擁有一個 PyTorch 模型，可以從 Python 例項化它，並且可以提供示例輸入張量來執行模型。

匯出和降低過程

將模型匯出並降低為 .pte 格式的過程通常包括以下步驟

1. 選擇要定位的後端。

2. 準備 PyTorch 模型，包括輸入和形狀規範。

3. 使用 torch.export.export 匯出模型。

4. 使用 to_edge_transform_and_lower 為目標後端最佳化模型。

5. 透過呼叫 to_executorch 並序列化輸出來建立 .pte 檔案。

量化 - 使用較低精度來減少推理時間和記憶體佔用的過程 - 通常也在這一階段進行。有關更多資訊，請參閱 量化概覽。

硬體後端

ExecuTorch 後端為特定硬體目標提供硬體加速。為了在目標硬體上實現最佳效能，ExecuTorch 在匯出和降低過程中會針對特定後端最佳化模型。這意味著生成的 .pte 檔案專門用於該特定硬體。為了部署到多個後端（例如 iOS 上的 Core ML 和 Android 上的 Arm CPU），通常會為每個後端生成一個專用的 .pte 檔案。

硬體後端的選擇取決於模型打算部署到的硬體。每個後端都有特定的硬體要求和模型支援級別。有關更多詳細資訊，請參閱每個硬體後端的文件。

作為 .pte 檔案建立過程的一部分，ExecuTorch 會識別給定後端支援的模型部分（分割槽）。這些部分會由後端預先處理以支援高效執行。模型中不受委託支援的部分（如果存在）將使用 CPU 上的可移植回退實現來執行。這允許在並非所有模型運算元都受後端支援時進行部分模型加速，但這可能會對效能產生負面影響。此外，可以按優先順序順序指定多個分割槽器。例如，這允許在 GPU 上不受支援的運算元透過 XNNPACK 在 CPU 上執行。

可用後端

下面列出了常用的硬體後端。對於移動裝置，請考慮在 Android 上使用 XNNPACK，在 iOS 上使用 XNNPACK 或 Core ML。要為特定後端建立 .pte 檔案，請將相應的分割槽器類傳遞給 to_edge_transform_and_lower。有關更多資訊，請參閱相應的後端文件以及下面的 匯出和降低 部分。

• XNNPACK（移動 CPU）

• Core ML (iOS)

• Metal Performance Shaders (iOS GPU)

• Vulkan (Android GPU)

• Qualcomm NPU

• MediaTek NPU

• Arm Ethos-U NPU

• Cadence DSP

模型準備

匯出過程接受一個標準的 PyTorch 模型，通常是 torch.nn.Module。這可以是自定義模型定義，也可以是來自現有來源（如 TorchVision 或 HuggingFace）的模型。有關降低 TorchVision 模型的示例，請參閱 ExecuTorch 入門。

模型匯出是透過 Python 完成的。這通常透過 Python 指令碼或互動式 Python Notebook（如 Jupyter 或 Colab）進行。下面的示例展示了一個簡單 PyTorch 模型的例項化和輸入。輸入被準備為 torch.Tensors 的元組，模型可以使用這些輸入執行。

import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.seq = torch.nn.Sequential(
            torch.nn.Conv2d(1, 8, 3),
            torch.nn.ReLU(),
            torch.nn.Conv2d(8, 16, 3),
            torch.nn.ReLU(),
            torch.nn.AdaptiveAvgPool2d((1,1))
       )
        self.linear = torch.nn.Linear(16, 10)
    
    def forward(self, x):
        y = self.seq(x)
        y = torch.flatten(y, 1)
        y = self.linear(y)
        return y

model = Model().eval()
inputs = (torch.randn(1,1,16,16),)
outputs = model(*inputs)
print(f"Model output: {outputs}")

請注意，模型使用 .eval() 設定為評估模式。除非執行裝置上訓練，否則模型應始終在評估模式下匯出。此模式將某些具有訓練特定行為的操作（例如 batch norm 或 dropout）配置為使用推理模式配置。

匯出和降低

要實際匯出和降低模型，請按順序呼叫 export、to_edge_transform_and_lower 和 to_executorch。這將生成一個 ExecuTorch 程式，可以將其序列化到檔案中。總而言之，使用 XNNPACK 委託（用於移動 CPU 效能）降低上面的示例模型可以按如下方式進行

import torch
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from executorch.exir import to_edge_transform_and_lower
from torch.export import Dim, export

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.seq = torch.nn.Sequential(
            torch.nn.Conv2d(1, 8, 3),
            torch.nn.ReLU(),
            torch.nn.Conv2d(8, 16, 3),
            torch.nn.ReLU(),
            torch.nn.AdaptiveAvgPool2d((1,1))
        )
        self.linear = torch.nn.Linear(16, 10)
    
    def forward(self, x):
        y = self.seq(x)
        y = torch.flatten(y, 1)
        y = self.linear(y)
        return y

model = Model()
inputs = (torch.randn(1,1,16,16),)
dynamic_shapes = {
    "x": {
        2: Dim("h", min=16, max=1024),
        3: Dim("w", min=16, max=1024),
    }
}

exported_program = export(model, inputs, dynamic_shapes=dynamic_shapes)
executorch_program = to_edge_transform_and_lower(
    exported_program,
    partitioner = [XnnpackPartitioner()]
).to_executorch()

with open("model.pte", "wb") as file:
    file.write(executorch_program.buffer)

這將生成一個 model.pte 檔案，該檔案可以在移動裝置上執行。

支援可變輸入大小（動態形狀）

PyTorch 匯出過程使用提供的示例輸入來跟蹤模型並推斷每一步的張量大小和型別。除非另有說明，否則匯出將假定固定的輸入大小等於示例輸入，並使用此資訊來最佳化模型。

許多模型需要支援可變輸入大小。為此，匯出接受一個 dynamic_shapes 引數，該引數告知編譯器哪些維度可以變化及其邊界。這採用巢狀字典的形式，其中鍵對應於輸入名稱，值指定每個輸入的邊界。

在示例模型中，輸入按照批次、通道、高度和寬度（NCHW）的標準慣例以 4 維張量形式提供。形狀為 [1, 3, 16, 16] 的輸入表示 1 個批次、3 個通道以及高度和寬度均為 16。

假設您的模型支援大小介於 16x16 和 1024x1024 之間的影像。形狀邊界可以按如下方式指定

dynamic_shapes = {
    "x": {
        2: Dim("h", min=16, max=1024),
        3: Dim("w", min=16, max=1024),
    }
}

ep = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)

在上面的示例中，"x" 對應於 Model.forward 中的引數名。鍵 2 和 3 對應於維度 2 和 3，即高度和寬度。由於沒有指定批次和通道維度，這些值根據示例輸入固定。

ExecuTorch 使用形狀邊界來最佳化模型並規劃模型執行的記憶體。因此，建議將維度上限設定得不要高於實際需要，因為上限越高，記憶體消耗越大。

對於更復雜的用例，動態形狀規範允許維度之間存在數學關係。有關動態形狀規範的更多資訊，請參閱 表達動態性。

測試模型

在整合執行時程式碼之前，通常會從 Python 測試匯出的模型。這可用於評估模型準確性和在移至目標裝置之前的健全性檢查。請注意，並非所有硬體後端都可以從 Python 中獲得，因為它們可能需要專用硬體才能執行。有關硬體要求和模擬器可用性的更多資訊，請參閱特定後端文件。此示例中使用的 XNNPACK 委託始終在主機上可用。

from executorch.runtime import Runtime

runtime = Runtime.get()

input_tensor = torch.randn(1, 3, 32, 32)
program = runtime.load_program("model.pte")
method = program.load_method("forward")
outputs = method.execute([input_tensor])

有關更多資訊，請參閱 執行時 API 參考。

後續步驟

PyTorch 和 ExecuTorch 的匯出和降低 API 提供了高度的可定製性，以滿足不同硬體和模型的需求。有關更多資訊，請參閱 torch.export 和 匯出 API 參考。

對於高階用例，請參閱以下內容

• 有關量化模型以減少推理時間和記憶體佔用的資訊，請參閱 量化概覽。

• 有關控制記憶體放置和規劃的資訊，請參閱 記憶體規劃。

• 有關編寫自定義編譯器 passes 的資訊，請參閱 自定義編譯器 Passes。

• 有關匯出生成的中間表示的資訊，請參閱 匯出 IR 規範。