在 C++ 中管理 Tensor 内存

作者： Anthony Shoumikhin

Tensor 是 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 保持有效。这使得代码维护困难且容易出错。用户一直在努力管理生命周期，许多人创建了自己的临时管理的 Tensor 抽象来将所有部分组合在一起，导致生态系统碎片化且不一致。

介绍 TensorPtr

为了缓解这些问题，ExecuTorch 提供了 TensorPtr，这是一个智能指针，用于管理 Tensor 数据及其动态元数据的生命周期。

使用 TensorPtr，用户不再需要单独担心元数据的生命周期。数据所有权取决于它是通过指针传递还是作为 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);

由于数据是以 vector 的形式提供的，因此数据现在由 Tensor 实例拥有。要创建非拥有的 TensorPtr，只需通过指针传递数据即可。type 是根据数据 vector（float）自动推导的。strides 和 dim_order 如果未明确指定为其他参数，则会根据 sizes 自动计算为默认值。

Module::forward() 中的 EValue 直接接受 TensorPtr，确保无缝集成。EValue 现在可以隐式构造为它可以容纳的任何类型的智能指针。这允许在将 TensorPtr 传递给 forward() 时隐式解引用 TensorPtr，并且 EValue 将持有 TensorPtr 指向的 Tensor。

API 概述

TensorPtr 字面意思是 std::shared_ptr<Tensor> 的别名，因此您可以轻松地使用它，而无需复制数据和元数据。每个 Tensor 实例可以拥有自己的数据，也可以引用外部数据。

创建 Tensor

有几种方法可以创建 TensorPtr。

创建标量 Tensor

您可以创建标量 Tensor，即零维 Tensor，或其中一个尺寸为零的 Tensor。

提供单个数据值

auto tensor = make_tensor_ptr(3.14);

生成的 Tensor 将包含单个 3.14 值，类型为 double，该类型会自动推导。

提供单个数据值并指定类型

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

现在，整数 42 将被转换为 float，Tensor 将包含单个 42 值，类型为 float。

从 Vector 拥有数据

当您提供大小和数据 Vector 时，TensorPtr 将拥有数据和大小的所有权。

提供数据 Vector

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

类型是从数据 Vector 自动推导为 ScalarType::Float。

提供数据 Vector 并指定类型

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

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

在此示例中，即使数据 Vector 包含整数，我们也将其标量类型指定为 Double。整数被转换为 double，并且新的数据 Vector 由 TensorPtr 拥有。由于在此示例中跳过了 sizes 参数，因此该 Tensor 是一个一维 Tensor，其大小等于数据 Vector 的长度。请注意，反向转换（从浮点类型到整型）是不允许的，因为它会丢失精度。同样，将其他类型转换为 Bool 也是不允许的。

将数据 Vector 提供为 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 Vector 必须足够大以容纳所有元素。

从原始指针获取非拥有数据

您可以创建引用现有数据但不拥有其所有权的 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,   // default dynamism
    [](void* ptr) { delete[] static_cast<double*>(ptr); });

TensorPtr 将在销毁时调用自定义析构函数，即当智能指针被重置且不再有对底层 Tensor 的引用时。

共享现有 Tensor

由于 TensorPtr 是一个 std::shared_ptr<Tensor>，您可以轻松创建共享现有 Tensor 的 TensorPtr。对共享数据的任何更改都会反映在所有共享相同数据的实例中。

共享现有 TensorPtr

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

现在 tensor 和 tensor_copy 指向相同的数据和元数据。

查看现有 Tensor

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

查看现有 Tensor

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

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

克隆 Tensor

要创建新的 TensorPtr，它拥有来自现有 Tensor 的数据副本

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 都将拥有数据的副本。

调整 Tensor 大小

TensorShapeDynamism 枚举指定了 Tensor 形状的可变性。

• STATIC： Tensor 的形状无法更改。

• DYNAMIC_BOUND： Tensor 的形状可以更改，但不能包含比创建时基于初始大小更多的元素。

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

调整 Tensor 大小时，您必须遵守其动态设置。仅允许调整具有 DYNAMIC 或 DYNAMIC_BOUND 形状的 Tensor，并且您不能将 DYNAMIC_BOUND Tensor 调整到比最初拥有的元素更多。

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 sizes are now {2, 2}
// Number of elements is 4 < initial 6

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

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

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

便捷助手

ExecuTorch 提供了一些便捷的助手函数来创建 Tensor。

使用 for_blob 和 from_blob 创建非拥有 Tensor

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

使用 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 在销毁时将调用自定义析构函数。

创建空 Tensor

empty() 创建一个具有指定大小的未初始化 Tensor。

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

empty_like() 创建一个具有与现有 TensorPtr 相同大小的未初始化 Tensor。

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

并且 empty_strided() 创建一个具有指定大小和步幅的未初始化 Tensor。

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

创建填充特定值的 Tensor

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

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 具有相同属性或具有自定义步幅的填充 Tensor。

创建随机 Tensor

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

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

randn() 创建一个填充有来自正态分布的随机值的 Tensor。

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

randint() 创建一个填充有指定范围（包含最小值，不包含最大值）内的随机整数的 Tensor。

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

创建标量 Tensor

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

auto tensor = scalar_tensor(3.14f);

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

关于 EValue 和生命周期管理的注意事项

Module 接口期望数据以 EValue 的形式提供，EValue 是一种变体类型，可以包含 Tensor 或其他标量类型。当您将 TensorPtr 传递给期望 EValue 的函数时，您可以解引用 TensorPtr 来获取底层 Tensor。

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

甚至可以是 EValues 的 vector，用于多个参数。

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::shared_ptr<at::Tensor>。这允许与 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 处理内存管理，也要确保任何非拥有的数据（例如，在使用 from_blob() 时）在 Tensor 使用期间保持有效。

• 使用便捷函数：利用助手函数来执行常见的 Tensor 创建模式，以编写更清晰、更易读的代码。

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

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

结论

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

通过提供与 PyTorch 的 ATen 库相似的接口，ExecuTorch 简化了新 API 的采用，使开发人员能够平滑过渡，无需陡峭的学习曲线。