在 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，这是一种智能指针，用于管理张量数据及其动态元数据的生命周期。

使用 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);

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

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

API 概述

TensorPtr 实际上是 std::shared_ptr<Tensor> 的别名，因此您可以轻松使用它而无需复制数据和元数据。每个 Tensor 实例可以拥有其数据或引用外部数据。

创建张量

有几种创建 TensorPtr 的方法。

创建标量张量

您可以创建标量张量，即零维张量或大小之一为零的张量。

提供单个数据值

auto tensor = make_tensor_ptr(3.14);

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

提供带有类型的单个数据值

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

现在，整数 42 将被强制转换为 float，并且张量将包含类型为 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。整数被强制转换为 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，该 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 的引用时，TensorPtr 将调用自定义删除器。

共享现有张量

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

共享现有 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 创建 TensorPtr，复制其属性并引用相同的数据。

查看现有张量

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

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

克隆张量

要创建一个新的 TensorPtr，该 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 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 提供了几个辅助函数来方便地创建张量。

使用 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

使用带有 Fluent 语法的 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 被销毁时，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 形式的数据，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，后者不拥有数据或元数据，而是使用原始指针引用它们。您需要确保 TensorPtr 在 EValue 使用期间保持有效。

当使用期望 EValue 的函数（如 set_input() 或 set_output()）时，这也适用。

与 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() 时）在张量使用期间保持有效。

• 使用便捷函数：利用辅助函数进行常见的张量创建模式，以编写更简洁易懂的代码。

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

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

结论

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

通过提供类似于 PyTorch 的 ATen 库的接口，ExecuTorch 简化了新 API 的采用，使开发人员可以在没有陡峭学习曲线的情况下进行过渡。