評價此頁面

張量基礎#

支援 PyTorch 的 ATen 張量庫是一個簡單的張量庫,它直接在 C++17 中公開了 Torch 中的張量操作。ATen 的 API 是從 PyTorch 使用的相同宣告中自動生成的,因此這兩個 API 將隨時間同步。

張量型別是動態解析的,因此 API 是通用的,不包含模板。也就是說,存在一個 Tensor 型別。它可以儲存 CPU 或 CUDA 張量,並且該張量可以包含 Doubles、Float、Ints 等型別的資料。這種設計使得無需對所有內容進行模板化即可輕鬆編寫通用程式碼。

請參閱 https://pytorch.com.tw/cppdocs/api/namespace_at.html#functions 以獲取提供的 API。摘錄

Tensor atan2(const Tensor & other) const;
Tensor & atan2_(const Tensor & other);
Tensor pow(Scalar exponent) const;
Tensor pow(const Tensor & exponent) const;
Tensor & pow_(Scalar exponent);
Tensor & pow_(const Tensor & exponent);
Tensor lerp(const Tensor & end, Scalar weight) const;
Tensor & lerp_(const Tensor & end, Scalar weight);
Tensor histc() const;
Tensor histc(int64_t bins) const;
Tensor histc(int64_t bins, Scalar min) const;
Tensor histc(int64_t bins, Scalar min, Scalar max) const;

還提供了原地操作,它們總是以 _ 為字尾,以表明它們將修改張量。

高效訪問張量元素#

使用張量範圍的操作時,動態分派的相對開銷非常小。然而,在某些情況下,特別是在您自己的核心中,需要高效的元素級訪問,此時元素級迴圈內的動態分派開銷會非常高。ATen 提供了訪問器,這些訪問器透過單次動態檢查張量的型別和維數來建立。然後,訪問器公開了一個用於高效訪問張量元素的 API。

訪問器是張量的臨時檢視。它們僅在所檢視的張量的生命週期內有效,因此應像迭代器一樣僅在函式內部區域性使用。

請注意,訪問器與核心函式內的 CUDA 張量不相容。相反,您必須使用打包訪問器,它的行為方式相同,但複製張量元資料而不是指向它。

因此,建議對 CPU 張量使用訪問器,對 CUDA 張量使用打包訪問器

CPU 訪問器#

torch::Tensor foo = torch::rand({12, 12});

// assert foo is 2-dimensional and holds floats.
auto foo_a = foo.accessor<float,2>();
float trace = 0;

for(int i = 0; i < foo_a.size(0); i++) {
  // use the accessor foo_a to get tensor data.
  trace += foo_a[i][i];
}

CUDA 訪問器#

__global__ void packed_accessor_kernel(
    torch::PackedTensorAccessor64<float, 2> foo,
    float* trace) {
  int i = threadIdx.x;
  gpuAtomicAdd(trace, foo[i][i]);
}

torch::Tensor foo = torch::rand({12, 12});

// assert foo is 2-dimensional and holds floats.
auto foo_a = foo.packed_accessor64<float,2>();
float trace = 0;

packed_accessor_kernel<<<1, 12>>>(foo_a, &trace);

除了 PackedTensorAccessor64packed_accessor64 之外,還有相應的 PackedTensorAccessor32packed_accessor32,它們使用 32 位整數進行索引。這在 CUDA 上會快得多,但可能導致索引計算中的溢位。

請注意,模板可以包含其他引數,例如指標限制和用於索引的整數型別。有關訪問器打包訪問器的詳細模板描述,請參閱文件。

使用外部建立的資料#

如果您已經在記憶體中分配了張量資料(CPU 或 CUDA),您可以在 ATen 中將該記憶體視為一個 Tensor

float data[] = { 1, 2, 3,
                 4, 5, 6 };
torch::Tensor f = torch::from_blob(data, {2, 3});

這些張量不能調整大小,因為 ATen 不擁有該記憶體,但它們在其他方面表現得像普通張量一樣。

標量和零維張量#

除了 Tensor 物件之外,ATen 還包括表示單個數字的 Scalar。與 Tensor 一樣,Scalar 是動態型別的,可以儲存 ATen 的任何數字型別。Scalar 可以從 C++ 數字型別隱式構造。需要 Scalar 是因為一些函式,例如 addmm,接受數字和張量,並期望這些數字與張量具有相同的動態型別。它們還用於 API 中,以指示函式總是返回 Scalar 值的位置,例如 sum

namespace torch {
Tensor addmm(Scalar beta, const Tensor & self,
             Scalar alpha, const Tensor & mat1,
             const Tensor & mat2);
Scalar sum(const Tensor & self);
} // namespace torch

// Usage.
torch::Tensor a = ...
torch::Tensor b = ...
torch::Tensor c = ...
torch::Tensor r = torch::addmm(1.0, a, .5, b, c);

除了 Scalar 之外,ATen 還允許 Tensor 物件是零維的。這些張量包含單個值,並且可以是更大 Tensor` 中單個元素的引用。它們可以在任何期望 Tensor 的地方使用。它們通常由像 select 這樣減少 Tensor` 維度的運算子建立。

torch::Tensor two = torch::rand({10, 20});
two[1][2] = 4;
// ^^^^^^ <- zero-dimensional Tensor