PyTorch 张量是如何实现的？

Posted 2023-02-16

【中文标题】PyTorch 张量是如何实现的？

【英文标题】：How are PyTorch's tensors implemented?

【发布时间】：2018-09-18 09:07:53

【问题描述】：

我正在用 Rust 构建我自己的 Tensor 类，并且我试图让它像 PyTorch 的实现一样。

以编程方式存储张量的最有效方法是什么，特别是在像 Rust 这样的强类型语言中？ 是否有任何资源可以很好地了解如何做到这一点？

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

张量的维度应该是动态的，这样我就可以拥有一个具有n 维度的张量。

【参考方案1】：

连续数组

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

唯一使用维度的地方是计算（许多）坐标和该数组中的偏移量之间的映射。例如，要获取项目[3, 1, 1]，您需要知道它是 3x3x3 矩阵、9x3x1 矩阵还是 27x1x1 矩阵 - 在所有情况下，“存储”的长度都是 27 个项目，但“坐标”的解释“会有所不同。如果使用从零开始的索引，计算很简单，但您需要知道每个维度的长度。

这确实意味着调整大小和类似操作可能需要复制整个数组，但没关系，您可以权衡这些（罕见）操作的性能以获得更常见操作的性能，例如顺序读取。

【参考方案2】：

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

每个张量都有一个关联的torch.Storage，它保存着它的数据。这 张量类提供存储的多维、跨步视图和 在其上定义数值运算。

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

Torch 支持 COO(rdinate) 格式的稀疏张量，可以 有效地存储和处理大多数的张量 元素为零。

如果您对不同稀疏格式之间的权衡感兴趣，那么从存储sparse matrices 的大量文献开始可能会很有用。许多用于矩阵的技术（实际上是 2 阶张量）转化为多维。 This masters dissertation 详细介绍了稀疏张量的具体实现。

***链接的摘要以及论文是

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

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

最后，我建议查看 sprs 以获取适用于 Rust 的良好稀疏矩阵库。在寻求扩展到多个维度时，这可能是一个不错的起点。