评价此页

张量创建 API#

本说明介绍如何在 PyTorch C++ API 中创建张量。它重点介绍了可用的工厂函数,这些函数根据某种算法填充新张量,并列出了可用于配置新张量的形状、数据类型、设备及其他属性的选项。

工厂函数#

工厂函数是一种生成新张量的函数。PyTorch(包括 Python 和 C++)中有许多可用的工厂函数,它们在返回新张量之前初始化张量的方式各不相同。所有工厂函数都遵循以下通用“模式”

torch::<function-name>(<function-specific-options>, <sizes>, <tensor-options>)

让我们分解此“模式”的各个部分

  1. <function-name> 是您想要调用的函数名称,

  2. <functions-specific-options> 是特定工厂函数接受的任何必需或可选参数,

  3. <sizes>IntArrayRef 类型的对象,用于指定结果张量的形状,

  4. <tensor-options>TensorOptions 的一个实例,用于配置结果张量的数据类型、设备、布局及其他属性。

选择工厂函数#

在撰写本文时,以下工厂函数可用(超链接指向相应的 Python 函数,因为它们的文档通常更详细——C++ 中的选项是相同的)

  • arange: 返回一个包含整数序列的张量,

  • empty: 返回一个带有未初始化值的张量,

  • eye: 返回一个单位矩阵,

  • full: 返回一个填充有单个值的张量,

  • linspace: 返回一个在某个区间内线性分布值的张量,

  • logspace: 返回一个在某个区间内对数分布值的张量,

  • ones: 返回一个填充有全一值的张量,

  • rand: 返回一个从 [0, 1) 上的均匀分布中抽取值填充的张量。

  • randint: 返回一个包含从某个区间随机抽取的整数的张量,

  • randn: 返回一个填充有从单位正态分布中抽取值的张量,

  • randperm: 返回一个填充有某个区间内的整数随机排列的张量,

  • zeros: 返回一个填充有全零值的张量。

指定尺寸#

按照张量填充方式的性质,不需要特定参数的函数可以直接通过尺寸调用。例如,以下行创建一个包含 5 个分量的向量,初始值全部设置为 1

torch::Tensor tensor = torch::ones(5);

如果我们想改用创建 3 x 5 矩阵,或者 2 x 3 x 4 张量呢?一般来说,IntArrayRef(工厂函数尺寸参数的类型)通过在大括号中指定每个维度的尺寸来构造。例如,对于具有两行三列的张量(在此情况下是矩阵),使用 {2, 3};对于三维张量,使用 {3, 4, 5};对于包含两个分量的一维张量,使用 {2}。在一维情况下,您可以省略大括号,只像上面那样传递单个整数。请注意,花括号只是构造 IntArrayRef 的一种方式。您也可以传递 std::vector<int64_t> 和其他一些类型。无论哪种方式,这意味着我们可以通过编写以下代码来构造一个填充有来自单位正态分布值的三维张量:

torch::Tensor tensor = torch::randn({3, 4, 5});
assert(tensor.sizes() == std::vector<int64_t>{3, 4, 5});

tensor.sizes() 返回一个 IntArrayRef,可以与 std::vector<int64_t> 进行比较,我们可以看到它包含了我们传递给张量的尺寸。您也可以编写 tensor.size(i) 来访问单个维度,这与 tensor.sizes()[i] 等价,但更推荐使用前者。

传递函数特定参数#

onesrandn 都不接受任何额外的参数来改变其行为。需要进一步配置的一个函数是 randint,它接受生成的整数值的上限,以及一个可选的下限(默认为零)。这里我们创建一个 5 x 5 的方阵,其中包含 0 到 10 之间的整数:

torch::Tensor tensor = torch::randint(/*high=*/10, {5, 5});

这里我们将下限提高到 3:

torch::Tensor tensor = torch::randint(/*low=*/3, /*high=*/10, {5, 5});

行内注释 /*low=*//*high=*/ 当然不是必需的,但它们像 Python 中的关键字参数一样有助于提高可读性。

提示

要点是尺寸总是紧随函数特定参数之后。

注意

有时函数完全不需要尺寸。例如,arange 返回的张量的尺寸完全由其函数特定参数(整数范围的下限和上限)指定。在这种情况下,该函数不接受 size 参数。

配置张量属性#

前面一节讨论了函数特定参数。函数特定参数只能更改填充张量的值,有时也会更改张量的尺寸。它们永远不会更改诸如正在创建的张量的数据类型(例如 float32int64),或者张量是否存储在 CPU 或 GPU 内存中。这些属性的指定留给了每个工厂函数的最后一个参数:一个 TensorOptions 对象,下面将进行讨论。

TensorOptions 是一个封装张量构造轴的类。这里的构造轴是指张量在构造之前(有时也可在构造之后)可以配置的特定属性。这些构造轴包括:

  • dtype(以前称为“标量类型”),它控制张量中存储元素的数据类型,

  • layout,可以是跨步式(密集)或稀疏式,

  • device,表示存储张量的计算设备(例如 CPU 或 CUDA GPU),

  • requires_grad 布尔值,用于启用或禁用张量的梯度记录,

如果您习惯使用 Python 中的 PyTorch,这些轴会听起来非常熟悉。目前这些轴的允许值包括:

  • 对于 dtypekUInt8, kInt8, kInt16, kInt32, kInt64, kFloat32kFloat64

  • 对于 layoutkStridedkSparse

  • 对于 devicekCPUkCUDA(接受可选的设备索引),

  • 对于 requires_gradtruefalse

提示

存在“Rust风格”的 dtype 缩写,例如 kF32 代替 kFloat32。请参阅此处查看完整列表。

TensorOptions 的一个实例存储了这些轴中每个轴的具体值。下面是一个创建 TensorOptions 对象的示例,该对象表示一个需要梯度计算的 64 位浮点型、跨步式张量,并且位于 CUDA 设备 1 上:

auto options =
  torch::TensorOptions()
    .dtype(torch::kFloat32)
    .layout(torch::kStrided)
    .device(torch::kCUDA, 1)
    .requires_grad(true);

注意我们如何使用 TensorOptions 的“构建器”风格方法来逐块构造对象。如果我们将此对象作为最后一个参数传递给工厂函数,则新创建的张量将具有这些属性:

torch::Tensor tensor = torch::full({3, 4}, /*value=*/123, options);

assert(tensor.dtype() == torch::kFloat32);
assert(tensor.layout() == torch::kStrided);
assert(tensor.device().type() == torch::kCUDA); // or device().is_cuda()
assert(tensor.device().index() == 1);
assert(tensor.requires_grad());

现在,您可能在想:我创建的每个新张量真的需要指定每个轴吗?幸运的是,答案是“不需要”,因为每个轴都有一个默认值。这些默认值是:

  • kFloat32 作为 dtype,

  • kStrided 作为 layout,

  • kCPU 作为 device,

  • false 作为 requires_grad

这意味着您在构造 TensorOptions 对象时省略的任何轴都将采用其默认值。例如,这是我们之前的 TensorOptions 对象,但 dtype 和 layout 已采用默认值:

auto options = torch::TensorOptions().device(torch::kCUDA, 1).requires_grad(true);

实际上,我们甚至可以省略所有轴来获得一个完全采用默认值的 TensorOptions 对象:

auto options = torch::TensorOptions(); // or `torch::TensorOptions options;`

这样做的一个好处是,我们刚才详细讨论的 TensorOptions 对象可以完全从任何张量工厂调用中省略:

// A 32-bit float, strided, CPU tensor that does not require a gradient.
torch::Tensor tensor = torch::randn({3, 4});
torch::Tensor range = torch::arange(5, 10);

但还有更便捷的方式:到目前为止,在此处介绍的 API 中,您可能注意到初始的 torch::TensorOptions() 写起来相当冗长。好消息是,对于每个构造轴(dtype、layout、device 和 requires_grad),torch:: 命名空间中都有一个自由函数,您可以将该轴的值传递给它。每个函数随后会返回一个已预先配置了该轴的 TensorOptions 对象,但仍允许通过上面所示的构建器风格方法进行进一步修改。例如:

torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32))

等价于

torch::ones(10, torch::dtype(torch::kFloat32))

此外,与其写成:

torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32).layout(torch::kStrided))

我们可以直接写成:

torch::ones(10, torch::dtype(torch::kFloat32).layout(torch::kStrided))

这为我们节省了不少打字工作。这意味着在实践中,您几乎(甚至从未)需要完整写出 torch::TensorOptions。取而代之的是使用 torch::dtype()torch::device()torch::layout()torch::requires_grad() 函数。

最后一个方便之处在于 TensorOptions 可以从单个值隐式构造。这意味着无论何时函数具有 TensorOptions 类型的参数(所有工厂函数都是如此),我们可以直接传递像 torch::kFloat32torch::kStrided 这样的值,而不是完整的对象。因此,当只需要更改相对于默认值的一个轴时,我们可以仅传递该值。例如,原来是:

torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32))

变成:

torch::ones(10, torch::dtype(torch::kFloat32))

并最终可以缩短为:

torch::ones(10, torch::kFloat32)

当然,使用这种简短语法无法修改 TensorOptions 实例的其他属性,但如果只需要更改一个属性,这非常实用。

总之,我们现在可以比较 TensorOptions 的默认值,以及使用自由函数创建 TensorOptions 的缩写 API,如何使 C++ 中的张量创建与 Python 中一样方便。比较 Python 中的这个调用:

torch.randn(3, 4, dtype=torch.float32, device=torch.device('cuda', 1), requires_grad=True)

与 C++ 中的等效调用:

torch::randn({3, 4}, torch::dtype(torch::kFloat32).device(torch::kCUDA, 1).requires_grad(true))

相当接近!

转换#

正如我们可以使用 TensorOptions 配置如何创建新张量一样,我们也可以使用 TensorOptions 将张量从一组属性转换为新的属性集。这种转换通常会创建一个新张量,并且不是原地操作。例如,如果有一个使用以下代码创建的 source_tensor

torch::Tensor source_tensor = torch::randn({2, 3}, torch::kInt64);

我们可以将其从 int64 转换为 float32

torch::Tensor float_tensor = source_tensor.to(torch::kFloat32);

注意

转换结果 float_tensor 是一个指向新内存的新张量,与源张量 source_tensor 无关。

然后我们可以将其从 CPU 内存移动到 GPU 内存:

torch::Tensor gpu_tensor = float_tensor.to(torch::kCUDA);

如果您有多个可用的 CUDA 设备,上述代码会将张量复制到默认的 CUDA 设备,您可以使用 torch::DeviceGuard 进行配置。如果没有设置 DeviceGuard,这将是 GPU 1。如果您想指定不同的 GPU 索引,可以将其传递给 Device 构造函数:

torch::Tensor gpu_two_tensor = float_tensor.to(torch::Device(torch::kCUDA, 1));

在 CPU 到 GPU 复制和反向操作的情况下,我们还可以通过将 /*non_blocking=*/false 作为最后一个参数传递给 to() 来将内存复制配置为异步

torch::Tensor async_cpu_tensor = gpu_tensor.to(torch::kCPU, /*non_blocking=*/true);

结论#

希望本说明能让您很好地理解如何使用 PyTorch C++ API 以惯用方式创建和转换张量。如果您有任何进一步的问题或建议,请使用我们的论坛GitHub issues 与我们联系。