在 C++ 中管理张量内存

作者： Anthony Shoumikhin

张量是 ExecuTorch 中的基本数据结构，表示用于神经网络和其他数值算法计算的多维数组。在 ExecuTorch 中，Tensor 类不拥有其元数据（大小、步幅、维度顺序）或数据，使运行时保持轻量级。用户负责提供所有这些内存缓冲区并确保元数据和数据在 Tensor 实例超出范围后仍然有效。虽然这种设计轻量级且灵活，尤其适用于小型嵌入式系统，但它给用户带来了很大的负担。但是，如果您的环境需要最少的动态分配、较小的二进制占用空间或有限的 C++ 标准库支持，则需要接受这种权衡并坚持使用常规的 Tensor 类型。

假设您正在使用 Module 接口，并且需要将 Tensor 传递给 forward() 方法。您需要至少分别声明和维护大小数组和数据，有时还需要步幅，这通常会导致以下模式

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

using namespace executorch::aten;
using namespace executorch::extension;

SizesType sizes[] = {2, 3};
DimOrderType dim_order[] = {0, 1};
StridesType strides[] = {3, 1};
float data[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
TensorImpl tensor_impl(
    ScalarType::Float,
    std::size(sizes),
    sizes,
    data,
    dim_order,
    strides);
// ...
module.forward(Tensor(&tensor_impl));

您必须确保 sizes、dim_order、strides 和 data 保持有效。这使得代码维护变得困难且容易出错。用户难以管理生命周期，许多用户创建了自己的临时管理张量抽象来将所有部分组合在一起，导致生态系统支离破碎且不一致。

引入 TensorPtr

为了缓解这些问题，ExecuTorch 通过新的 张量扩展 提供了 TensorPtr 和 TensorImplPtr 来管理张量及其实现的生命周期。它们本质上是智能指针（分别为 std::unique_ptr<Tensor> 和 std::shared_ptr<TensorImpl>），负责处理张量数据及其动态元数据的内存管理。

现在，用户不再需要单独担心元数据生命周期。数据所有权根据它是通过指针传递还是作为 std::vector 移动到 TensorPtr 中来确定。所有内容都捆绑在一个地方并自动管理，使您能够专注于实际计算。

以下是使用方法

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

using namespace executorch::extension;

auto tensor = make_tensor_ptr(
    {2, 3},                                // sizes
    {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}); // data
// ...
module.forward(tensor);

数据现在由张量实例拥有，因为它作为向量提供。要创建一个非拥有 TensorPtr，只需通过指针传递数据即可。type 会根据数据向量 (float) 自动推断。如果未明确指定为额外参数，则 strides 和 dim_order 会根据 sizes 自动计算为默认值。

EValue 在 Module::forward() 中直接接受 TensorPtr，确保无缝集成。EValue 现在可以使用指向它可以保存的任何类型的智能指针隐式构造，因此 TensorPtr 会隐式取消引用，并将 EValue 持有的 TensorPtr 指向的 Tensor 传递给 forward()。

API 概述

新的 API 围绕两个主要的智能指针展开

• TensorPtr：std::unique_ptr 管理 Tensor 对象。由于每个 Tensor 实例都是唯一的，因此 TensorPtr 确保独占所有权。

• TensorImplPtr：std::shared_ptr 管理 TensorImpl 对象。多个 Tensor 实例可以共享相同的 TensorImpl，因此 TensorImplPtr 使用共享所有权。

创建张量

有几种方法可以创建 TensorPtr。

创建标量张量

您可以创建一个标量张量，即一个维度为零或其中一个大小为零的张量。

提供单个数据值

auto tensor = make_tensor_ptr(3.14);

生成的张量将包含一个类型为 double 的单个值 3.14，该类型会自动推断。

使用类型提供单个数据值

auto tensor = make_tensor_ptr(42, ScalarType::Float);

现在整数 42 将转换为浮点数，并且张量将包含一个类型为 float 的单个值 42。

拥有数据向量

当您提供尺寸和数据向量时，TensorPtr 将同时拥有数据和尺寸。

提供数据向量

auto tensor = make_tensor_ptr(
    {2, 3},                                 // sizes
    {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f});  // data (float)

数据类型会根据数据向量自动推断为 ScalarType::Float。

使用类型提供数据向量

如果您提供了一种类型的数据，但指定了不同的标量类型，则数据将被转换为指定的类型。

auto tensor = make_tensor_ptr(
    {1, 2, 3, 4, 5, 6},          // data (int)
    ScalarType::Double);         // double scalar type

在这个例子中，即使数据向量包含整数，我们也指定标量类型为 Double。整数将被转换为双精度浮点数，新的数据向量由 TensorPtr 拥有。此示例中跳过了 sizes 参数，因此使用输入数据向量的大小。请注意，我们禁止反向转换，即浮点类型转换为整型类型，因为这会导致精度损失。类似地，不允许将其他类型转换为 Bool。

将数据向量作为 std::vector<uint8_t> 提供

您还可以将原始数据作为 std::vector<uint8_t> 提供，并指定尺寸和标量类型。数据将根据提供的类型重新解释。

std::vector<uint8_t> data = /* raw data */;
auto tensor = make_tensor_ptr(
    {2, 3},                 // sizes
    std::move(data),        // data as uint8_t vector
    ScalarType::Int);       // int scalar type

data 向量必须足够大，以根据提供的尺寸和标量类型容纳所有元素。

不拥有原始数据指针

您可以创建一个引用现有数据但不拥有其所有权的 TensorPtr。

提供原始数据

float data[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
auto tensor = make_tensor_ptr(
    {2, 3},              // sizes
    data,                // raw data pointer
    ScalarType::Float);  // float scalar type

TensorPtr 不拥有数据，您必须确保 data 保持有效。

使用自定义删除器提供原始数据

如果您希望 TensorPtr 管理数据的生命周期，您可以提供一个自定义删除器。

auto* data = new double[6]{1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
auto tensor = make_tensor_ptr(
    {2, 3},                               // sizes
    data,                                 // data pointer
    ScalarType::Double,                   // double scalar type
    TensorShapeDynamism::DYNAMIC_BOUND,   // some default dynamism
    [](void* ptr) { delete[] static_cast<double*>(ptr); });

TensorPtr 将在销毁时调用自定义删除器，即当智能指针被重置并且不再有对底层 TensorImplPtr 的引用时。

共享现有张量

您可以通过包装现有的 TensorImplPtr 创建一个 TensorPtr，后者可以使用与 TensorPtr 相同的 API 集合创建。对 TensorImplPtr 或任何共享相同 TensorImplPtr 的 TensorPtr 进行的任何更改都会反映在所有中。

共享现有 TensorImplPtr

auto tensor_impl = make_tensor_impl_ptr(
    {2, 3},
    {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f});
auto tensor = make_tensor_ptr(tensor_impl);
auto tensor_copy = make_tensor_ptr(tensor_impl);

tensor 和 tensor_copy 都共享底层的 TensorImplPtr，反映数据中的更改，但不反映元数据中的更改。

此外，您可以创建一个新的 TensorPtr，它与现有的 TensorPtr 共享相同的 TensorImplPtr。

共享现有 TensorPtr

auto tensor_copy = make_tensor_ptr(tensor);

查看现有张量

您可以从现有的 Tensor 创建一个 TensorPtr，复制其属性并引用相同的数据。

查看现有张量

Tensor original_tensor = /* some existing tensor */;
auto tensor = make_tensor_ptr(original_tensor);

现在新创建的 TensorPtr 引用与原始张量相同的数据，但拥有自己的元数据副本，因此可以以不同的方式解释或“查看”数据，但对数据的任何修改也将反映到原始 Tensor 中。

克隆张量

要创建一个新的 TensorPtr，它拥有来自现有张量的数据副本

Tensor original_tensor = /* some existing tensor */;
auto tensor = clone_tensor_ptr(original_tensor);

新创建的 TensorPtr 拥有自己的数据副本，因此可以独立地修改和管理它。同样，您可以创建现有 TensorPtr 的克隆。

auto original_tensor = make_tensor_ptr();
auto tensor = clone_tensor_ptr(original_tensor);

请注意，无论原始 TensorPtr 是否拥有数据，新创建的 TensorPtr 都将拥有数据副本。

调整张量大小

TensorShapeDynamism 枚举指定了张量形状的可变性

• STATIC：张量的形状无法更改。

• DYNAMIC_BOUND：张量的形状可以更改，但永远不能包含比根据初始尺寸在创建时具有的更多元素。

• DYNAMIC：张量的形状可以任意更改。请注意，当前 DYNAMIC 是 DYNAMIC_BOUND 的别名。

调整张量大小时，必须遵守其动态设置。仅允许对具有 DYNAMIC 或 DYNAMIC_BOUND 形状的张量进行调整大小，并且您不能将 DYNAMIC_BOUND 张量调整为包含比其最初具有的更多元素。

auto tensor = make_tensor_ptr(
    {2, 3},                      // sizes
    {1, 2, 3, 4, 5, 6},          // data
    ScalarType::Int,
    TensorShapeDynamism::DYNAMIC_BOUND);
// Initial sizes: {2, 3}
// Number of elements: 6

resize_tensor_ptr(tensor, {2, 2});
// The tensor's sizes are now {2, 2}
// Number of elements is 4 < initial 6

resize_tensor_ptr(tensor, {1, 3});
// The tensor's sizes are now {1, 3}
// Number of elements is 3 < initial 6

resize_tensor_ptr(tensor, {3, 2});
// The tensor's sizes are now {3, 2}
// Number of elements is 6 == initial 6

resize_tensor_ptr(tensor, {6, 1});
// The tensor's sizes are now {6, 1}
// Number of elements is 6 == initial 6

便捷助手

ExecuTorch 提供了一些辅助函数，可以方便地创建张量。

使用 for_blob 和 from_blob 创建非拥有张量

这些助手允许您创建不拥有数据的张量。

使用 from_blob()

float data[] = {1.0f, 2.0f, 3.0f};
auto tensor = from_blob(
    data,                // data pointer
    {3},                 // sizes
    ScalarType::Float);  // float scalar type

使用流畅语法使用 for_blob()

double data[] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
auto tensor = for_blob(data, {2, 3}, ScalarType::Double)
                  .strides({3, 1})
                  .dynamism(TensorShapeDynamism::STATIC)
                  .make_tensor_ptr();

使用 from_blob() 的自定义删除器

int* data = new int[3]{1, 2, 3};
auto tensor = from_blob(
    data,             // data pointer
    {3},              // sizes
    ScalarType::Int,  // int scalar type
    [](void* ptr) { delete[] static_cast<int*>(ptr); });

TensorPtr 将在销毁时调用自定义删除器。

创建空张量

empty() 创建一个使用指定尺寸的未初始化张量。

auto tensor = empty({2, 3});

empty_like() 创建一个与现有 TensorPtr 尺寸相同的未初始化张量。

TensorPtr original_tensor = /* some existing tensor */;
auto tensor = empty_like(original_tensor);

而 empty_strided() 创建一个使用指定尺寸和步长的未初始化张量。

auto tensor = empty_strided({2, 3}, {3, 1});

创建填充特定值的张量

full()、zeros() 和 ones() 分别创建填充提供值、零或一的张量。

auto tensor_full = full({2, 3}, 42.0f);
auto tensor_zeros = zeros({2, 3});
auto tensor_ones = ones({3, 4});

与 empty() 类似，还有额外的辅助函数 full_like()、full_strided()、zeros_like()、zeros_strided()、ones_like() 和 ones_strided() 用于创建与现有 TensorPtr 属性相同或具有自定义步长的填充张量。

创建随机张量

rand() 创建一个填充 0 到 1 之间随机值的张量。

auto tensor_rand = rand({2, 3});

randn() 创建一个填充正态分布随机值的张量。

auto tensor_randn = randn({2, 3});

randint() 创建一个填充指定最小值（包含）和最大值（不包含）之间的随机整数的张量。

auto tensor_randint = randint(0, 10, {2, 3});

创建标量张量

除了使用单个数据值的 make_tensor_ptr() 之外，您还可以使用 scalar_tensor() 创建标量张量。

auto tensor = scalar_tensor(3.14f);

请注意，scalar_tensor() 函数期望值为 Scalar 类型。在 ExecuTorch 中，Scalar 可以表示 bool、int 或浮点类型，但不能表示 Half 或 BFloat16 等类型，对于这些类型，您需要使用 make_tensor_ptr() 跳过 Scalar 类型。

关于 EValue 和生命周期管理的说明

Module 接口期望数据采用 EValue 的形式，这是一种变体类型，可以保存 Tensor 或其他标量类型。当您将 TensorPtr 传递给期望 EValue 的函数时，您可以取消引用 TensorPtr 以获取底层的 Tensor。

TensorPtr tensor = /* create a TensorPtr */
//...
module.forward(tensor);

甚至可以是多个参数的 EValues 向量。

TensorPtr tensor = /* create a TensorPtr */
TensorPtr tensor2 = /* create another TensorPtr */
//...
module.forward({tensor, tensor2});

但是，请注意：EValue 不会保存来自 TensorPtr 的动态数据和元数据。它仅仅保存一个普通的 Tensor，该 Tensor 不拥有数据或元数据，而是使用原始指针引用它们。您需要确保 TensorPtr 在 EValue 使用期间保持有效。

这同样适用于使用诸如 set_input() 或 set_output() 等函数的情况，这些函数期望 EValue 作为输入。

与 ATen 的互操作性

如果您的代码在启用预处理器标志 USE_ATEN_LIB 的情况下编译，则所有 TensorPtr API 将在后台使用 at:: API。例如，TensorPtr 将变为 std::unique_ptr<at::Tensor>，而 TensorImplPtr 将变为 c10::intrusive_ptr<at::TensorImpl>。这使得与 PyTorch ATen 库无缝集成成为可能。

API 等效表

这是一个将 TensorPtr 创建函数与其对应的 ATen API 进行匹配的表格

ATen	ExecuTorch
at::tensor(data, type)	make_tensor_ptr(data, type)
at::tensor(data, type).reshape(sizes)	make_tensor_ptr(sizes, data, type)
tensor.clone()	clone_tensor_ptr(tensor)
tensor.resize_(new_sizes)	resize_tensor_ptr(tensor, new_sizes)
at::scalar_tensor(value)	scalar_tensor(value)
at::from_blob(data, sizes, type)	from_blob(data, sizes, type)
at::empty(sizes)	empty(sizes)
at::empty_like(tensor)	empty_like(tensor)
at::empty_strided(sizes, strides)	empty_strided(sizes, strides)
at::full(sizes, value)	full(sizes, value)
at::full_like(tensor, value)	full_like(tensor, value)
at::full_strided(sizes, strides, value)	full_strided(sizes, strides, value)
at::zeros(sizes)	zeros(sizes)
at::zeros_like(tensor)	zeros_like(tensor)
at::zeros_strided(sizes, strides)	zeros_strided(sizes, strides)
at::ones(sizes)	ones(sizes)
at::ones_like(tensor)	ones_like(tensor)
at::ones_strided(sizes, strides)	ones_strided(sizes, strides)
at::rand(sizes)	rand(sizes)
at::rand_like(tensor)	rand_like(tensor)
at::randn(sizes)	randn(sizes)
at::randn_like(tensor)	randn_like(tensor)
at::randint(low, high, sizes)	randint(low, high, sizes)
at::randint_like(tensor, low, high)	randint_like(tensor, low, high)

最佳实践

• 仔细管理生命周期：即使 TensorPtr 和 TensorImplPtr 处理内存管理，您仍然需要确保任何非拥有数据（例如，当使用 from_blob() 时）在张量使用期间保持有效。

• 使用便捷函数：利用提供的辅助函数来处理常见的张量创建模式，以编写更简洁、更易读的代码。

• 了解数据所有权：了解您的张量是否拥有其数据或引用外部数据，以避免意外的副作用或内存泄漏。

• 确保 TensorPtr 的生命周期长于 EValue：当将张量传递给期望 EValue 的模块时，请确保 TensorPtr 在 EValue 使用期间保持有效。

• 理解标量类型：创建张量时，尤其是在类型之间进行转换时，请注意标量类型。

结论

ExecuTorch 中的 TensorPtr 和 TensorImplPtr 通过将数据和动态元数据捆绑到智能指针中来简化张量内存管理。这种设计消除了用户管理多个数据片段的需要，并确保了更安全、更易于维护的代码。

通过提供类似于 PyTorch 的 ATen 库的接口，ExecuTorch 使开发人员更容易采用新的 API，而无需陡峭的学习曲线。