目录

C++中运行ExecuTorch模型教程

作者: Jacob Szwejbka

在本教程中,我们将介绍如何使用更详细、更低级别的API在C++中运行ExecuTorch模型:准备MemoryManager,设置输入,执行模型并检索输出。然而,如果你正在寻找一个开箱即用的更简单的接口,请考虑尝试模块扩展教程

关于ExecuTorch运行时的高级概述,请参阅运行时概述,更多关于每个API的详细文档请参阅运行时API参考这里是一个功能齐全的C++模型运行器版本,《设置ExecuTorch》文档展示了如何构建和运行它。

预备知识

您需要一个ExecuTorch模型来跟随学习。我们将使用从导出到ExecuTorch教程生成的模型SimpleConv

模型加载

运行模型的第一步是加载模型。ExecuTorch 使用一种称为 DataLoader 的抽象来处理获取 .pte 文件数据的具体细节,然后 Program 表示已加载的状态。

用户可以定义自己的DataLoader以适应其特定系统的需求。在本教程中,我们将使用FileDataLoader, 但您可以查看示例数据加载器实现以了解ExecuTorch项目提供的其他选项。

对于 FileDataLoader,我们只需要向构造函数提供一个文件路径。

using executorch::aten::Tensor;
using executorch::aten::TensorImpl;
using executorch::extension::FileDataLoader;
using executorch::extension::MallocMemoryAllocator;
using executorch::runtime::Error;
using executorch::runtime::EValue;
using executorch::runtime::HierarchicalAllocator;
using executorch::runtime::MemoryManager;
using executorch::runtime::Method;
using executorch::runtime::MethodMeta;
using executorch::runtime::Program;
using executorch::runtime::Result;
using executorch::runtime::Span;

Result<FileDataLoader> loader =
        FileDataLoader::from("/tmp/model.pte");
assert(loader.ok());

Result<Program> program = Program::load(&loader.get());
assert(program.ok());

设置MemoryManager

接下来我们将设置 MemoryManager

ExecuTorch 的一个原则是让用户能够控制运行时所使用的内存来源。今天(2023年底),用户需要提供 2 种不同的分配器:

  • 方法分配器:一个 MemoryAllocator,用于在 Method 加载时分配运行时结构。诸如张量元数据、内部指令链以及其他运行时状态等信息都来自这里。

  • 计划内存:一个 HierarchicalAllocator,包含 1 个或多个内存区域,内部可变张量数据缓冲区被放置在此处。在 Method 加载时,内部张量的数据指针会被分配到各个偏移量位置内。这些偏移量的位置和区域的大小是在内存规划阶段提前确定的。

在这个示例中,我们将从 Program 动态获取计划内存区域的大小,但对于无堆环境,用户可以提前从 Program 获取此信息并静态分配区域。我们还将使用基于 malloc 的分配器作为方法分配器。

// Method names map back to Python nn.Module method names. Most users will only
// have the singular method "forward".
const char* method_name = "forward";

// MethodMeta is a lightweight structure that lets us gather metadata
// information about a specific method. In this case we are looking to get the
// required size of the memory planned buffers for the method "forward".
Result<MethodMeta> method_meta = program->method_meta(method_name);
assert(method_meta.ok());

std::vector<std::unique_ptr<uint8_t[]>> planned_buffers; // Owns the Memory
std::vector<Span<uint8_t>> planned_arenas; // Passed to the allocator

size_t num_memory_planned_buffers = method_meta->num_memory_planned_buffers();

// It is possible to have multiple layers in our memory hierarchy; for example,
// SRAM and DRAM.
for (size_t id = 0; id < num_memory_planned_buffers; ++id) {
  // .get() will always succeed because id < num_memory_planned_buffers.
  size_t buffer_size =
      static_cast<size_t>(method_meta->memory_planned_buffer_size(id).get());
  planned_buffers.push_back(std::make_unique<uint8_t[]>(buffer_size));
  planned_arenas.push_back({planned_buffers.back().get(), buffer_size});
}
HierarchicalAllocator planned_memory(
    {planned_arenas.data(), planned_arenas.size()});

// Version of MemoryAllocator that uses malloc to handle allocations rather then
// a fixed buffer.
MallocMemoryAllocator method_allocator;

// Assemble all of the allocators into the MemoryManager that the Executor will
// use.
MemoryManager memory_manager(&method_allocator, &planned_memory);

加载一个方法

在ExecuTorch中,我们以方法粒度从Program加载和初始化。许多程序只会有一个方法‘forward’。load_method是进行初始化的地方,包括设置张量元数据、初始化委托等。

Result<Method> method = program->load_method(method_name);
assert(method.ok());

设置输入

现在我们已经有了我们的方法,我们需要在进行推理之前设置其输入。在这种情况下,我们知道我们的模型接受一个大小为 (1, 3, 256, 256) 的浮点张量。

根据你的模型内存规划方式,规划好的内存可能包含也可能不包含用于输入和输出的缓冲区空间。

如果输出数据没有经过内存规划,用户将需要使用 'set_output_data_ptr' 设置输出数据指针。在这种情况下,我们将假设模型已经通过内存规划处理了输入和输出。

// Create our input tensor.
float data[1 * 3 * 256 * 256];
Tensor::SizesType sizes[] = {1, 3, 256, 256};
Tensor::DimOrderType dim_order = {0, 1, 2, 3};
TensorImpl impl(
    ScalarType::Float, // dtype
    4, // number of dimensions
    sizes,
    data,
    dim_order);
Tensor t(&impl);

// Implicitly casts t to EValue
Error set_input_error = method->set_input(t, 0);
assert(set_input_error == Error::Ok);

执行推理

现在我们的方法已经加载并且输入已经设置,我们可以进行推理。我们通过调用 execute 来执行此操作。

Error execute_error = method->execute();
assert(execute_error == Error::Ok);

获取输出

一旦我们的推理完成,我们可以获取输出。我们知道模型只会返回一个输出张量。这里的一个潜在问题是,我们得到的输出是 Method 所拥有。用户在对其进行任何修改之前,应确保克隆输出,或者如果需要让输出的生命周期独立于 Method,则应这样做。

EValue output = method->get_output(0);
assert(output.isTensor());

结论

本教程演示了如何使用低级运行时API运行ExecuTorch模型,这些API提供了对内存管理和执行的细粒度控制。然而,对于大多数用例,我们建议使用模块API,它提供了更简化的体验而不牺牲灵活性。有关更多详情,请查看模块扩展教程

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源