將模型降階為委派 (Delegate)¶

目標對象：對應用委派 (Delegate) 來加速運行時程式感興趣的 ML 工程師。

後端委派 (Backend delegation) 是一種後端處理和執行 PyTorch 程式的入口點，旨在利用特定後端和硬體的效能和效率優勢，同時為 PyTorch 使用者提供接近 PyTorch 運行時的體驗。後端委派通常由 ExecuTorch 或供應商提供。在程式中使用委派的方法是透過標準的入口點 to_backend。

前端介面¶

將程式委派給後端有三種流程：

將整個模組降級 (Lower) 到後端。這適用於測試後端和前處理階段。
將整個模組降級到後端，並將其與另一個模組組合。這適用於重複使用從其他流程匯出的降級模組。
根據分割器 (Partitioner) 降級模組的部分節點。這適用於降級包含可降級和不可降級節點的模型，並且是最簡化的流程。

流程 1：降級整個模組¶

此流程從具有邊緣方言 (Edge Dialect) 表示法的追蹤圖模組開始。要降級它，我們呼叫以下函式，該函式會傳回 LoweredBackendModule (有關此函式的更多文件，請參閱匯出 API 參考)

# defined in backend_api.py
def to_backend(
    backend_id: str,
    edge_program: ExportedProgram,
    compile_spec: List[CompileSpec],
) -> LoweredBackendModule:

在此函式中，會呼叫後端的 preprocess() 函式，該函式會產生已編譯的 blob，該 blob 將發送到 FlatBuffer 二進位檔。降級後的模組可以直接捕獲，也可以放回父模組中進行捕獲。最終，捕獲的模組將在 FlatBuffer 的模型中序列化，並且可以由運行時載入。

以下是此流程的範例：

from executorch.exir.backend.backend_api import to_backend
import executorch.exir as exir
import torch
from torch.export import export
from executorch.exir import to_edge

# The submodule runs in a specific backend. In this example,  `BackendWithCompilerDemo` backend
class LowerableSubModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return torch.sin(x)

# Convert the lowerable module to Edge IR Representation
to_be_lowered = LowerableSubModel()
example_input = (torch.ones(1), )
to_be_lowered_exir_submodule = to_edge(export(to_be_lowered, example_input))

# Import the backend implementation
from executorch.exir.backend.test.backend_with_compiler_demo import (
    BackendWithCompilerDemo,
)
lowered_module = to_backend('BackendWithCompilerDemo', to_be_lowered_exir_submodule.exported_program(), [])

我們可以透過直接運行以下程式碼，將程式序列化為 FlatBuffer 格式：

# Save the flatbuffer to a local file
save_path = "delegate.pte"
with open(save_path, "wb") as f:
    f.write(lowered_module.buffer())

流程 2：降級整個模組並組合¶

或者，在流程 1 之後，我們可以將此降級模組與另一個模組組合：

# This submodule runs in executor runtime
class NonLowerableSubModel(torch.nn.Module):
    def __init__(self, bias):
        super().__init__()
        self.bias = bias

    def forward(self, a, b):
        return torch.add(torch.add(a, b), self.bias)


# The composite module, including lower part and non-lowerpart
class CompositeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.non_lowerable = NonLowerableSubModel(torch.ones(1) * 0.3)
        self.lowerable = lowered_module

    def forward(self, x):
        a = self.lowerable(x)
        b = self.lowerable(a)
        ret = self.non_lowerable(a, b)
        return a, b, ret

composite_model = CompositeModel()
model_inputs = (torch.ones(1), )
exec_prog = to_edge(export(composite_model, model_inputs)).to_executorch()

# Save the flatbuffer to a local file
save_path = "delegate.pte"
with open(save_path, "wb") as f:
    f.write(exec_prog.buffer)

流程 3：分割¶

第三個流程也從具有邊緣方言表示法的追蹤圖模組開始。要降級此圖模組中的某些節點，我們可以使用重載的 to_backend 函式。

def to_backend(
    edge_program: ExportedProgram,
    partitioner: Partitioner,
) -> ExportedProgram:

此函式接收一個 Partitioner，該分割器會將標籤新增到所有要降級的節點。它將傳回一個 partition_tags 字典，將標籤對應到後端名稱和模組編譯規格。然後，將使用流程 1 的流程對標記的節點進行分割並降級到其對應的後端。可用的輔助分割器記錄在此處。這些降級後的模組將插入到頂層模組中並序列化。