作为委托降低模型¶

受众：机器学习工程师，他们有兴趣通过委托来加速程序的运行时性能。

后端委托是后端处理和执行 PyTorch 程序的入口点，旨在利用专用后端和硬件在性能与效率方面的优势，同时为 PyTorch 用户提供接近 PyTorch 运行时的体验。后端委托通常由 ExecuTorch 或供应商提供。在程序中利用委托的方式是通过标准入口点 to_backend。

前端接口¶

有三种将程序委托给后端的流程：

将整个模块降低到后端。这有助于测试后端和预处理阶段。
将整个模块降低到后端并与另一个模块组合。这对于重用从其他流程导出的已降低模块很有帮助。
模块根据分区器划分的下部部分。这对于包含可降低和不可降低节点的模型进行降低非常有用，并且是流程最简洁的方式。

流程 1: 降低整个模块¶

此流程从具有Edge Dialect表示的跟踪图模块开始。要降低它，我们调用以下函数，该函数返回LoweredBackendModule（有关此函数的更多文档，请参见< a data-trkey="3">导出API参考）

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

在此函数中，会调用后端的 preprocess() 函数，该函数生成一个编译后的 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: 分割¶

第三个流程也从一个带有Edge Dialect表示的 traced graph 模块开始。为了降低此图模块中的某些节点，我们可以使用重载的to_backend函数。

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

此函数接收一个Partitioner，这将为所有需要降低的节点添加标签。它将返回一个partition_tags字典，该字典将标签映射到后端名称和模块编译规范。然后，使用Flow 1的过程将标记的节点分区并降低到其映射的后端。可用的帮助程序分区器文档请参阅这里。这些已降低的模块将被插入到顶级模块中并序列化。

以下是一个流程示例：

import executorch.exir as exir
from executorch.exir.backend.backend_api import to_backend
from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo
from executorch.exir.program import (
    EdgeProgramManager,
    to_edge,
)
from torch.export import export
import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, y):
        x = x + y
        x = x * y
        x = x - y
        x = x / y
        x = x * y
        x = x + y
        return x

model = Model()
model_inputs = (torch.randn(1, 3), torch.randn(1, 3))

core_aten_ep = export(model, model_inputs)
edge: EdgeProgramManager = to_edge(core_aten_ep)
edge = edge.to_backend(AddMulPartitionerDemo())
exec_prog = edge.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)

运行时¶

在拥有委托程序之后，要运行带有后端的模型，我们需要注册后端。根据委托的实现方式，后端可以作为全局变量的一部分进行注册，也可以显式地在主函数内部进行注册。

如果在全局变量初始化期间进行注册，只要它是静态链接的，后端将被注册。用户只需将库作为依赖项包含进去即可。
如果供应商提供了注册后端的 API，用户需要将库作为依赖项包含进来，并在主函数中调用供应商提供的 API 显式注册后端。