在 C++ 中使用模块扩展运行 ExecuTorch 模型¶

作者： Anthony Shoumikhin

在在 C++ 中运行 ExecuTorch 模型教程中，我们探索了用于运行导出模型的较低级别的 ExecuTorch API。虽然这些 API 提供了零开销、极大的灵活性和控制力，但对于常规使用来说，它们可能显得冗长而复杂。为了简化这一点，并使其类似于 Python 中 PyTorch 的 eager 模式，我们引入了常规 ExecuTorch 运行时 API 之上的 Module 外观 API。Module API 提供了相同的灵活性，但默认使用常用组件，如 DataLoader 和 MemoryAllocator，从而隐藏了大多数复杂的细节。

示例¶

让我们看看如何使用 Module 和 TensorPtr API 运行从导出到 ExecuTorch 教程生成的 SimpleConv 模型

#include <executorch/extension/module/module.h>
#include <executorch/extension/tensor/tensor.h>

using namespace ::executorch::extension;

// Create a Module.
Module module("/path/to/model.pte");

// Wrap the input data with a Tensor.
float input[1 * 3 * 256 * 256];
auto tensor = from_blob(input, {1, 3, 256, 256});

// Perform an inference.
const auto result = module.forward(tensor);

// Check for success or failure.
if (result.ok()) {
  // Retrieve the output data.
  const auto output = result->at(0).toTensor().const_data_ptr<float>();
}

现在的代码简化为创建一个 Module 并在其上调用 forward()，无需额外的设置。让我们仔细看看这些和其他 Module API，以便更好地理解内部工作原理。

API¶

创建模块¶

创建 Module 对象是一个快速操作，不涉及大量的处理时间或内存分配。除非通过专用 API 显式请求，否则 Program 和 Method 的实际加载会在首次推理时延迟发生。

Module module("/path/to/model.pte");

强制加载方法¶

要随时强制加载 Module（以及由此产生的底层 ExecuTorch Program），请使用 load() 函数

const auto error = module.load();

assert(module.is_loaded());

要强制加载特定的 Method，请调用 load_method() 函数

const auto error = module.load_method("forward");

assert(module.is_method_loaded("forward"));

您还可以使用便捷函数加载 forward 方法

const auto error = module.load_forward();

assert(module.is_method_loaded("forward"));

注意： Program 会在任何 Method 加载之前自动加载。如果之前的尝试成功，则后续加载它们的尝试无效。

查询元数据¶

使用 method_names() 函数获取 Module 包含的方法名称集

const auto method_names = module.method_names();

if (method_names.ok()) {
  assert(method_names->count("forward"));
}

注意： 首次调用 method_names() 时，它将强制加载 Program。

要内省关于特定方法的各种元数据，请使用 method_meta() 函数，该函数返回 MethodMeta 结构

const auto method_meta = module.method_meta("forward");

if (method_meta.ok()) {
  assert(method_meta->name() == "forward");
  assert(method_meta->num_inputs() > 1);

  const auto input_meta = method_meta->input_tensor_meta(0);
  if (input_meta.ok()) {
    assert(input_meta->scalar_type() == ScalarType::Float);
  }

  const auto output_meta = method_meta->output_tensor_meta(0);
  if (output_meta.ok()) {
    assert(output_meta->sizes().size() == 1);
  }
}

注意： 首次调用 method_meta() 时，它也会强制加载 Method。

执行推理¶

假设提前知道 Program 的方法名称及其输入格式，您可以使用 execute() 函数直接按名称运行方法

const auto result = module.execute("forward", tensor);

对于标准的 forward() 方法，上述操作可以简化为

const auto result = module.forward(tensor);

注意： 首次调用 execute() 或 forward() 时，它们将加载 Program 和 Method。因此，首次推理将花费更长的时间，因为模型是延迟加载的，并且除非之前已显式加载，否则会为执行做好准备。

设置输入和输出¶

您可以使用以下 API 为方法设置单独的输入和输出值。

设置输入¶

输入可以是任何 EValue，包括张量、标量、列表和其他支持的类型。要为方法设置特定的输入值

module.set_input("forward", input_value, input_index);

input_value 是一个 EValue，表示您要设置的输入。
input_index 是要设置的输入的从零开始的索引。

例如，要设置第一个输入张量

module.set_input("forward", tensor_value, 0);

您还可以一次设置多个输入

std::vector<runtime::EValue> inputs = {input1, input2, input3};
module.set_inputs("forward", inputs);

注意： 您可以跳过 forward() 方法的方法名称参数。

通过预先设置所有输入，您可以执行推理而无需传递任何参数

const auto result = module.forward();

或者仅设置然后部分传递输入

// Set the second input ahead of time.
module.set_input(input_value_1, 1);

// Execute the method, providing the first input at call time.
const auto result = module.forward(input_value_0);

注意： 预设输入存储在 Module 中，并且可以在后续执行中多次重复使用。

如果您不再需要预设输入，请不要忘记清除或重置输入，方法是将它们设置为默认构造的 EValue

module.set_input(runtime::EValue(), 1);

设置输出¶

运行时只能设置 Tensor 类型的输出，并且它们在模型导出时不得进行内存规划。内存规划的张量在模型导出期间预先分配，无法替换。

要为特定方法设置输出张量

module.set_output("forward", output_tensor, output_index);

output_tensor 是一个 EValue，包含您要设置为输出的张量。
output_index 是要设置的输出的从零开始的索引。

注意： 确保您设置的输出张量与方法输出的预期形状和数据类型匹配。

您可以跳过 forward() 的方法名称和第一个输出的索引

module.set_output(output_tensor);

注意： 与输入一样，预设输出也存储在 Module 中，并且可以在后续执行中多次重复使用。

结果和错误类型¶

大多数 ExecuTorch API 返回 Result 或 Error 类型

Error 是一个 C++ 枚举，包含有效的错误代码。默认值为 Error::Ok，表示成功。
Result 可以保存 Error（如果操作失败），或有效负载，例如包装 Tensor 的 EValue（如果成功）。要检查 Result 是否有效，请调用 ok()。要检索 Error，请使用 error()，要获取数据，请使用 get() 或解引用运算符（如 * 和 ->）。

分析模块¶

使用ExecuTorch Dump跟踪模型执行。创建一个 ETDumpGen 实例并将其传递给 Module 构造函数。执行方法后，将 ETDump 数据保存到文件中以进行进一步分析

#include <fstream>
#include <memory>

#include <executorch/extension/module/module.h>
#include <executorch/devtools/etdump/etdump_flatcc.h>

using namespace ::executorch::extension;

Module module("/path/to/model.pte", Module::LoadMode::MmapUseMlock, std::make_unique<ETDumpGen>());

// Execute a method, e.g., module.forward(...); or module.execute("my_method", ...);

if (auto* etdump = dynamic_cast<ETDumpGen*>(module.event_tracer())) {
  const auto trace = etdump->get_etdump_data();

  if (trace.buf && trace.size > 0) {
    std::unique_ptr<void, decltype(&free)> guard(trace.buf, free);
    std::ofstream file("/path/to/trace.etdump", std::ios::binary);

    if (file) {
      file.write(static_cast<const char*>(trace.buf), trace.size);
    }
  }
}

结论¶

Module API 为在 C++ 中运行 ExecuTorch 模型提供了简化的接口，与 PyTorch eager 模式的体验非常相似。通过抽象出较低级别运行时 API 的复杂性，开发者可以专注于模型执行，而无需担心底层细节。