ExecuTorch-rs

License

ExecuTorch 的绑定 - 面向 PyTorch 的移动、嵌入式和边缘设备上的端侧 AI。

提供了一个高级 Rust API，用于使用 ExecuTorch 库，具体是 C++ API，在移动、嵌入式和边缘设备上执行 PyTorch 模型。 PyTorch 模型在 Python 中创建和导出，然后使用 ExecuTorch 库在设备端加载和执行。

以下示例在 Python 中创建一个简单模型，导出它，然后在 Rust 中执行它：

在 Python 中创建模型并导出：

import torch
from torch.export import export
from executorch.exir import to_edge_transform_and_lower


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

    def forward(self, x: torch.Tensor, y: torch.Tensor):
        return x + y


model = Add()
exported_program = export(model, (torch.ones(1), torch.ones(1)))
executorch_program = to_edge_transform_and_lower(exported_program).to_executorch()
with open("model.pte", "wb") as file:
    file.write(executorch_program.buffer)

在 Rust 中执行模型：

use executorch::evalue::{EValue, IntoEValue};
use executorch::module::Module;
use executorch::tensor_ptr;
use ndarray::array;

let mut module = Module::new("model.pte");

let (tensor1, tensor2) = (tensor_ptr![1.0_f32], tensor_ptr![1.0_f32]);
let inputs = [tensor1.into_evalue(), tensor2.into_evalue()];

let outputs = module.forward(&inputs).unwrap();
let [output]: [EValue; 1] = outputs.try_into().expect("not a single output");
let output = output.as_tensor().into_typed::<f32>();

println!("Output tensor computed: {:?}", output);
assert_eq!(array![2.0], output.as_array());

参见 example/hello_world 获取完整示例。

构建

要使用该库，您必须自行编译 C++ executorch 库，因为有许多配置决定了支持哪些模块、后端和操作。有关更多信息，请参阅 executorch-sys crate。当前支持的 Cpp executorch 版本是 1.3.1。

下表显示了哪个版本的 Rust crate 与哪个 C++ ExecuTorch 版本兼容：

Rust crate	C++ 库
0.11.x	1.3.1
0.10.x	1.2.0
0.9.x	1.1.0
0.8.x	1.0.1
0.7.x	0.7.0
0.6.x	0.6.0
0.5.x	0.5.0
0.4.x	0.4.0
0.3.x	0.3.0
0.2.x	0.3.0
0.1.x	0.2.1

Cargo 特性

data-loader: 在 data_loader 模块中包含用于加载数据的额外结构体。如果没有此特性，唯一可用的数据加载器是 BufferDataLoader。需要 libextension_data_loader.a 静态库，使用 EXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON 编译 C++ executorch。
module: 包含 module API，这是一个用于加载和执行 PyTorch 模型的高级 API。它是低级别 Program API 的替代方案，后者更适合嵌入式系统。需要 libextension_module_static.a 静态库，使用 EXECUTORCH_BUILD_EXTENSION_MODULE=ON 编译 C++ executorch。还包含 std、data-loader 和 flat-tensor 特性。
tensor-ptr: 包含 tensor::TensorPtr 结构体，一种用于张量的智能指针，它管理张量对象的生命周期，同时管理数据缓冲区以及额外元数据的生命周期。需要 extension_tensor.a 静态库，使用 EXECUTORCH_BUILD_EXTENSION_TENSOR=ON 编译 C++ executorch。还包含 std 特性。
flat-tensor: 包含 FlatTensorDataMap 结构体，可以读取带有外部张量的模型的 .ptd 文件。需要 libextension_flat_tensor.a 静态库，使用 EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON 编译 C++ executorch。
etdump: 包含 ETDumpGen 结构体，它是 EventTracer 的一种实现，用于调试和分析。需要 libetdump.a 静态库，使用 EXECUTORCH_BUILD_DEVTOOLS=ON 和 EXECUTORCH_ENABLE_EVENT_TRACER=ON 编译 C++ executorch。此外，需要 flatcc（或 flatcc_d）库，可在 {CMAKE_DIR}/third-party/flatcc_ep/lib/ 获取，并且应由用户进行链接。
ndarray: executorch 张量与 ndarray 数组之间的转换。添加对 ndarray crate 的依赖。此特性默认启用。
f16: 添加对 half crate 的依赖，该 crate 提供功能完整的 f16 和 bf16 类型。如果未启用此功能，这两种类型仅可通过简单的转换与 u16 相互转换。请注意，这仅影响输入/输出张量，内部计算始终具备处理此类标量的能力。
num-complex: 添加对 num-complex crate 的依赖，该 crate 提供了功能完备的复数类型。如果未启用此功能，复数将作为一个具有两个公开字段且无任何操作的简单结构体可用。请注意，这仅影响输入/输出张量，内部计算始终具备处理此类标量的能力。
std: 启用标准库。此功能默认启用，但可以禁用以在 no_std 环境中构建 executorch 。参见 examples/no_std 示例。同时包含 alloc 功能。注意：no_std 仍处于开发中，参见 https://github.com/pytorch/executorch/issues/4561
alloc: 启用内存分配。当此功能被禁用时，所有需要内存分配的方法将不会被编译。此功能由 std 功能启用，而该功能默认启用。可以在不启用 std 功能的情况下启用此功能，此时内存分配将使用 alloc crate 完成，该 crate 要求设置全局分配器。

默认情况下，std 和 ndarray 功能是启用的。