TensorFlow/PyTorch中张量(Tensor)的底层存储方式

---

@[toc]

1 张量（Tensor）基本概念回顾

张量（Tensor）其实就是多维数组，类似于NumPy里面的np.array。
这里的维度，更准确的讲法应该叫阶（rank），这是为了跟向量（vector）的维度区分开的。vector其实就是rank为1的张量，我们说一个vector是n维的其实是说它有n个分量（标量）。而如果张量的维度（阶）是n维的，并不是说它有n个标量分量，而是说在表示这个张量时需要用n个坐标轴。每个轴上都可以有多个分量。为了表示每个轴上的分量个数，引入形状（shape）的概念。
例如，一个三阶的张量，可以理解为是一个立方体。假设它的shape是$[3,2,5]$,那么下图就是一个具体的例子（下面两张图的来源均为https://www.tensorflow.org/guide/tensor）：

图0.1 一个rank为3的tensor的例子

而一个四阶张量的例子则用下图表示：

图0.2 一个rank为4的tensor的例子

TensorFlow/PyTorch中使用比较多的tensor的阶为4，shape为$[Batch, Height, Weight, Features]$.

n阶张量的排列规律如下图所示（图片来源https://syborg.dev/posts/understanding-tensors-in-pytorch）：

图0.3 1-6阶tensor的shape规律

可以将规律总结为：从shape列表的最右边往左遍历，最开始三个阶按照“下-右-里”的顺序排列，然后打包成一个group，再将整个group按照“下-右-里”的顺序排列，满三次后再打包成一个group，如此往复循环。。

2 tensor在计算机内存中的存储方式

前面提到的rank和shape，都是数学上的定义，实际在内存中是一维存储的。
排列规律为：

1. 假设rank为n，即shape列表的size为n.
2. 初始位置为原点
3. 将shape[n-1]方向上的元素按照这个方向上的坐标递增的顺序进行线性排列。
4. shape[n-1]方向上的元素排完后再将shape[n-2]上的坐标递增，继续排shape[n-1]方向上的元素，直到shape[n-2]方向上的坐标到了最大值，并且排完了这个坐标上的shape[n-1]方向的元素，此时shape[n-2]方向上坐标归零。
5. 将shape[n-3]上的坐标递增，继续3,4，直到shape[n-3]方向上的坐标和shape[n-2]方向上的坐标到了最大值，并且排完了这个坐标上的shape[n-1]方向的元素，此时shape[n-3]和shape[n-2]方向上坐标归零
6. 对shape[n-i] (i=4,5,…,n)上的坐标递增,并重复前面的过程。

一般情况下，假设高阶tensor的各个下标存在数组coordinates中, 各个阶的维数存在数组shape中，则转换为一维向量的下标index过程如下：

int index=0;
for(int i=0;i<coordinates.size(); i++){
    int tmp=coordinates[shape.size()-1-i];
    for(int j=0;j<i;j++){
        tmp=tmp*shape[shape.size()-1-j];
    }
    index=index+tmp;
}

当然，还有一种基于动态规划的更高效的计算方法：

int index = coordinates[0];
for(int i=1; i<coordinates.size(); i++){
index = index * shape[i] + coordinates[i];
}