Fork me on GitHub

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过程如下:

1
2
3
4
5
6
7
8
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;
}

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

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