在 C++ 中管理 Tensor 記憶體

作者： Anthony Shoumikhin

在 ExecuTorch 中，張量 (Tensor) 是基礎的資料結構，代表多維陣列，用於神經網路和其他數值演算法的計算。在 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) 自動推導的。如果沒有作為額外參數明確指定，則會根據 sizes 自動計算 strides 和 dim_order 的預設值。

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，該指標引用現有資料而不取得所有權。

提供原始資料

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 是一個 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 中。

複製 Tensors

要建立一個新的 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 都將擁有資料的副本。

調整 Tensors 大小

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 提供了幾個輔助函數，可以方便地建立 tensors。

使用 for_blob 和 from_blob 建立非擁有的 Tensors

這些輔助函數允許你建立不擁有資料的 tensors。

使用 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() 的自訂 Deleter

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 將在銷毀時呼叫自訂 deleter。

建立空 Tensors

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});

建立填充特定值的 Tensors

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 相同的屬性或具有自訂步幅的填充 tensors。

建立隨機 Tensors

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});

建立純量 Tensors

除了使用單個資料值的 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 形式的資料，這是一種變體類型，可以容納 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 的動態資料和 metadata。它僅持有常規的 Tensor，該 Tensor 不擁有資料或 metadata，而是使用原始指標 (raw pointers) 參考它們。您需要確保 TensorPtr 在 EValue 使用期間保持有效。

當使用如 set_input() 或 set_output() 等需要 EValue 的函式時，也適用此規則。

與 ATen 的互操作性

如果您的程式碼編譯時啟用了 preprocessor flag 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：當將 tensors 傳遞給需要 EValue 的模組時，請確保 TensorPtr 在 EValue 使用期間保持有效。

結論

ExecuTorch 中的 TensorPtr 透過將資料和動態 metadata 捆綁到智慧指標中，簡化了 tensor 記憶體管理。這種設計消除了使用者管理多個資料片段的需求，並確保更安全、更易於維護的程式碼。

透過提供與 PyTorch 的 ATen 函式庫相似的介面，ExecuTorch 簡化了新 API 的採用，使開發人員能夠在沒有陡峭學習曲線的情況下進行轉換。