PyTorch框架Tensor底层原理

神经网络，可以看做是一个非常复杂的公式的代码实现。计算机没法理解抽象的公式，需要我们翻译翻译。例如 $y = 10x + 3$这个公式，我们想让计算机能用它来做计算，我们需要用编程语言做个转换，C++实现如下：

template<class T>
T linear(T x) {
  T a = 10;
  T b = 3;
  return a * x + b;
}

这个公式最终是平面中一系列的点所组成的一条直线。

当来到神经网络这个非常复杂的公式，最终你会在一个超空间里表现出一个我们无法想象的平面。我们每一条数据，都可以看成是一个超空间里面的点，所有数据在一起应该会组成一个超平面。神经网络的目的就是企图找到这个超平面。搭建模型的过程，就是提供一个有潜力拟合这个超平面的一部分的过程；训练的过程就是让这个公式尝试去拟合这个超平面的一部分的过程。通过训练，合适的模型最终会拟合一条曲线、平面、超平面一部分。为什么说是一部分呢，因为数据是有限的，我们所找到的这个公式也许最终也只是在我们所给出的数据范围内能够拟合这个超平面。超平面是固有的，公式是我们试图对这个超平面的复现。

神经网络经过学习，可以将一种信息转换成另信息的另一种描述。例如将一张猫的图片转换为“猫”这个文字。因为信息在神经网络中都是以浮点数来表现的，我们需要把我们的信息都转化到浮点数上来（定点也行，不过原理都一样）。这样我们的每一条数据就能表现为超平面上的一个点了。

信息转换成浮点数值之后，还需要有一个数据结构来存储，并且计算的中间结果也是一些浮点数，也需要一个数据结构来存储。

通常，经过转换的信息和计算的中间结果都会表现为一个多维数组，例如可以用一个三维数组表示图片、四维数组表示视频等。

虽然Python中的List可以表示多维数组，但是神经网络中一般不会使用List去存储数据，原因主要有以下几点：
首先，Python List是一个对象的集合，它可以存储任何对象，即便存储的是浮点数值，每一个浮点数值都经过对应Python对象的封装，添加了额外的空间去保存引用计数等信息，这会造成很多空间了浪费，并且离散的分布在内存中，不利于进行优化；
其次，Python List是没有定义有对其所表示的向量、矩阵等的点乘等操作，这些操作在神经网络中是基本的操作；
再有，对List的操作需要通过Python解析器来完成，相对于直接执行机器指令而言，这种方法速度是非常慢的。

这些缺点对于需要进行密集计算的神经网络来说，是无法忍受的。因此，需要一个便捷、高效的数据结构来存储计算过程中的中间数据等多维数组。不同的框架可能称呼不一样，例如在NumPy中称为ndarray。而在PyTorch中这个专门的数据结构称为Tensor。

Tensor（张量），可能在物理等其他领域表示的意义略有差别，但在计算机中的表现形式其实就是一个多维数组，就类似于一维数组称被称为向量、二维数组被称为矩阵一样。

Tensor的底层是用C/C++实现的，它不仅用于存储数据，还定义了对数据的操作。抛开这些不说，它与Python List最大的区别就是它使用一个连续的内存区域去存储数据，并且这些数据并未经过封装。Python List 和Tensor的区别如图2所示：

在图2中，左边表示Python List，可以看到每个数值类型都被封装成了一个PyObject对象，每个对象是独立分配内存、离散的存储在内存中；右边表示Tensor，Tensor中的数值统一的保存在一块连续的内存中，而且是不经过封装的。

你以为你已经看到Tensor的内部原理了？不，你没有。

真正管理存储这数据的内存区域的，是类Storage的实例，这个Storage的实例通过一个一维数组来存储数据。你没看错，是一维数组。不管外在表现为多少维的数组，都是存储在一个一维数组中。而怎么让这个一维数组看起来像多维数组，就是Tensor完成的。其内部实现关系如图3所示。

Storage类中有一个指针指向存储数据的一位数组，而Tensor通过对Storage进行封装，使得在外部看来数据是多维的。
多个Tensor的实例可以指向同一个Storage，例如下面的代码：

import torch
a = torch.tensor([[4,1,5],[3,2,1]])
print(a.storage())
ar = a.reshape((3,2))
print(ar.storage())
at = a.transpose(1,0)
print(at.storage())

输出如下：

 4
 1
 5
 3
 2
 1
[torch.LongStorage of size 6]
 4
 1
 5
 3
 2
 1
[torch.LongStorage of size 6]
 4
 1
 5
 3
 2
 1
[torch.LongStorage of size 6]

他们看起来是不同的Tensor，但是指向的却是同一个Storage，如果你该表其中一个的内容，另一个Tensor也会随之改变，例如，我们将上面例子中ar的第二行第一个元素的值变成100，看看会发生什么：

ar[1,1] = 100
print(a.storage())
print(ar.storage())
print(at.storage())

输出：

 4
 1
 5
 100
 2
 1
[torch.LongStorage of size 6]
 4
 1
 5
 100
 2
 1
[torch.LongStorage of size 6]
 4
 1
 5
 100
 2
 1
[torch.LongStorage of size 6]

可以看到，a和at中的元素也被改变了。上面的例子可以用图4形象的表示出来。

Tensor是怎么做到在不对底层存储数据的一维数组做任何改变的情况下，让他看起来是多个不同的多维数组的呢？答案是Tensor有三个帮手：Size、storage offset和strides。

Size是一个torch.Size对象，它保存着一个指示每个维度上有多少个元素的列表，它控制着每个维度的取值范围。例如下面的例子中有个3x3的Tensor，通过获取它的一个2x2切分，可以看到除了Size值改变了，其他的都没有变，引用的还是同一个Storage；

a = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print(a)
print(a.storage())
print("size is {}".format(a.size()))
print("stride is {}".format(a.stride()))
print("storage offset is {}".format(a.storage_offset()))

Output:

tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])
 1
 2
 3
 4
 5
 6
 7
 8
 9
[torch.LongStorage of size 9]
size is torch.Size([3, 3])
stride is (3, 1)
storage offset is 0

Slice:

sa = a[:2, :2]
print(sa)
print(sa.storage())
print("size is {}".format(sa.size()))
print("stride is {}".format(sa.stride()))
print("storage offset is {}".format(sa.storage_offset()))

Output：

tensor([[1, 2],
        [4, 5]])
 1
 2
 3
 4
 5
 6
 7
 8
 9
[torch.LongStorage of size 9]
size is torch.Size([2, 2])
stride is (3, 1)
storage offset is 0

storage offset是一个指向该Tensor元素开始的Storage索引，因为可能有些Tensor只使用了Storage的一部分，它控制着每个Tensor的起始位置。
strides是一个元组，与用于表示获取一个Storage中的一个元素在每个维度上需要跳过个元素，它控制着如何得到某个元素在Storage上的索引。
他们的关系如图5所示。

通过这三个属性，结合给出的维度信息，就可以计算出任意元素在Storage上的索引，计算公式如下：

$n_{offset} = \sum_{k=0}^{N-1}s_kn_k + {storage\_offset}$
其中： $s_k$表示stride； $n_k$表示对应的维度。
例如上面3x3例子中矩阵的例子中，如果我们想获取第二行第二列的元素，我们给出的下标索引为a[1,1]，而a.stride()的值为(3, 1)，那么Storage中的索引值为：
$n_{offset} = 1\cdot3 + 1\cdot1 + 0 = 4$
也就是Storage的第五个元素，即5。类似的，我们可以算出a[2,1] = 7, a[0, 2] = 2，Storage中对应的值分别是8和3。Size并没有直接参与公式计算，它用来控制我们给出的下标，防止我们越界访问。

现在我们知道通过stride和storage offset可以得到 $n_{offset}$，那么stride又是怎么来的呢？同样有公式：
$s_k = itemsize \prod_{j=k+1}^{N-1}d_j$

例如，我们又一个1x4x4x3的Tensor，那么根据公式，我们可以算出 $s_0 = 4\cdot4\cdot3 = 48$，同理可算出 $s_1 s_2 s_3$，最终得到stride = (48, 12, 3, 1)，那么到底是不是呢？我们做个小实验：

a = torch.randn(1,4,4,3)
print(a.stride())

输出为：

(48, 12, 3, 1)

和我们算的一样。

而transpose操作的结果，实际上就是改变stride元组中元素的顺序，例如：

at = a.transpose(1,3)
print(at.stride())

输出为：

(48, 1, 3, 12)

可以看到，就是1和12位置调换了，最终导致同一个下标算出的索引值改变。

总结

PyTorch中数据全部保存在一个由Storage对象维护的一位数组中（更确切的说是一块连续的内存区域），而Tensor只不过是对这个一维数组的一个视角。Tensor通过Size, storage offset, strides这三个属性的值的不同组合，可以让同一个Storage一维数组看起来像多个不同的多维数组。

首发于个人微信公众号TensorBoy。微信扫描上方二维码或者微信搜索TensorBoy并关注，及时获取更多最新文章！
C++ | Python | 推理引擎 | AI框架源码，有一起玩耍的么？

References

[1] https://docs.scipy.org/doc/numpy/reference/arrays.ndarray.html#internal-memory-layout-of-an-ndarray
[2] Deep-Learning-with-PyTorch.pdf
[3] https://github.com/pytorch/pytorch