张量基础知识¶
支持 PyTorch 的 ATen 张量库是一个简单的张量库,它直接在 C++17 中公开了 Torch 中的张量操作。ATen 的 API 是从 PyTorch 使用的相同声明自动生成的,因此这两个 API 将随着时间的推移而相互跟踪。
张量类型是动态解析的,因此 API 是通用的,不包括模板。也就是说,有一个 Tensor
类型。它可以保存 CPU 或 CUDA 张量,并且张量可以包含双精度浮点数、单精度浮点数、整数等。这种设计使编写通用代码变得容易,而无需对所有内容进行模板化。
请参阅 https://pytorch.ac.cn/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);
除了 PackedTensorAccessor64
和 packed_accessor64
之外,还有相应的 PackedTensorAccessor32
和 packed_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
,它表示单个数字。与张量一样,标量是动态类型的,可以保存 ATen 的任何一种数字类型。标量可以从 C++ 数字类型隐式构造。标量是必要的,因为某些函数(如 addmm
)需要数字和张量,并期望这些数字与张量具有相同的动态类型。它们也用于 API 中,以指示函数将始终返回标量值的位置,例如 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