将模型作为委托降低¶
受众:对将委托应用于运行时加速程序感兴趣的 ML 工程师。
后端委托是后端处理和执行 PyTorch 程序的入口点,以利用专门后端和硬件的性能和效率优势,同时仍为 PyTorch 用户提供接近 PyTorch 运行时的体验。后端委托通常由 ExecuTorch 或供应商提供。在程序中利用委托的方法是通过标准入口点 to_backend
。
前端接口¶
将程序委托给后端有三种流程
将整个模块降低到后端。这对测试后端和预处理阶段很有用。
将整个模块降低到后端并将其与另一个模块组合。这对重用从其他流程导出的已降低模块很有用。
根据分区器降低模块的某些部分。这对降低包含可降低节点和不可降低节点的模型很有用,也是最简化的流程。
流程 1:降低整个模块¶
此流程从具有边缘方言表示的跟踪图模块开始。要降低它,我们调用以下函数,该函数返回一个 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 的过程将标记的节点分区并降低到其映射的后端。可用的辅助分区器在 此处有说明。这些已降低的模块将被插入到顶级模块中并被序列化。
以下示例说明了此流程
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 在主函数中显式注册后端。