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