PyTorch 的张量是如何实现的?

How are PyTorch's tensors implemented?

我正在用 Rust 构建我自己的 Tensor class,我正试图让它像 PyTorch 的实现一样。

以编程方式存储张量的最有效方法是什么,但是,特别是在像 Rust 这样的强类型语言中? 是否有任何资源可以提供良好的深入了解这是如何完成的?

我目前正在构建一个连续的数组,因此,给定 3 x 3 x 3 的维度,我的数组中将只有 3^3 个元素,这将代表张量。但是,这确实使数组的一些数学运算和操作变得更加困难。

张量的维度应该是动态的,这样我就可以得到一个 n 维度的张量。

连续数组

存储此类数据的常用方法是在单个数组中,该数组在内存中布置为单个连续块。更具体地说,一个 3x3x3 张量将简单地存储为一个包含 27 个值的数组,一个接一个。

唯一使用维度的地方是计算(许多)坐标和该数组内的偏移量之间的映射。例如,要获取项目 [3, 1, 1],您需要知道它是 3x3x3 矩阵、9x3x1 矩阵还是 27x1x1 矩阵 - 在所有情况下,"storage" 的长度都是 27 个项目,但是"coordinates" 的解释会有所不同。如果使用zero-based索引,计算很简单,但需要知道每个维度的长度。

这确实意味着调整大小和类似操作可能需要复制整个数组,但没关系,您可以牺牲那些(罕见的)操作的性能来获得更常见操作的性能,例如顺序读取。

PyTorch 默认以密集格式存储其张量。根据the docs,

Each tensor has an associated torch.Storage, which holds its data. The tensor class provides multi-dimensional, strided view of a storage and defines numeric operations on it.

正如您想象的那样,当存储具有许多零的大张量时,存储所有值的内存使用效率非常低。因此,PyTorch 也提供 sparse tensors。来自该页面:

Torch supports sparse tensors in COO(rdinate) format, which can efficiently store and process tensors for which the majority of elements are zeros.

如果您对不同稀疏格式之间的权衡感兴趣,从有关存储的大量文献入手可能会很有用 sparse matrices. Many of the techniques used for matrices (which are effectively rank 2 tensors) translate to multiple dimensions. This masters dissertation 特别详细地介绍了稀疏张量的实现。

维基百科的摘要 link 以及论文

  • 易于指定和修改的格式通常计算效率低下,因此稀疏结构通常使用一种规范构建但存储在另一种规范中

  • 选择符合您预期用途的存储格式很重要。例如,矩阵-矩阵乘法需要高效的列访问,而矩阵-向量乘法需要高效的行访问。

最后,我建议查看 sprs 以寻找适用于 Rust 的良好稀疏矩阵库。当希望扩展到多个维度时,它可能是一个很好的起点。