原版英文鏈接:Edward Z. Yang's PyTorch internals : Inside 245-5D
Tensor
張量(Tensor)是PyTorch中最核心的數(shù)據(jù)結(jié)構(gòu),包含了相關(guān)元數(shù)據(jù)的描述,包括張量的大小(size)、所存儲單元數(shù)據(jù)的類型(dtype)、存儲位置(device)、步幅(stride)、以及布局(layout)。
Stride
張量是一種數(shù)學(xué)概念,實際應(yīng)用中需要存儲到計算機物理存儲空間中,為了提升訪問效率,需要在物理空間上連續(xù)存儲(通常按行(row)來展開)。為了在連續(xù)存儲的同時保留張量的大小維度信息,需要使用步幅作為額外信息描述。每個維度上的步幅描述了在此維度上兩個邏輯位置相鄰的元素在實際物理存儲位置上的間隔,從而實現(xiàn)了從邏輯地址到物理存儲地址的映射。
例如:
對于張量a=Tensor([ [1, 2], [3, 4] ]),其維度大小為[2, 2],對應(yīng)在各維度的步幅為[2, 1],分別表示在零維(row)的步幅為2,第一維(column)的步幅為1。對于張量b=a[1, 0],其在張量a中的邏輯位置為(1, 0),對應(yīng)的物理存儲位置則為:1*2+0*1=2,即各維度上的邏輯索引與步幅的乘積之和。
對于同一個張量,如果希望獲取某一行的張量數(shù)據(jù),例如c=a[1, :],此時需記錄offset=1*2=2,而在另外一個緯度上則繼承了張量a的步幅,因此對于張量c,其size為[2],步幅則為[1], offset為2。需要注意的是,張量c和張量a是共享同一塊物理存儲,即任何一個修改,都會影響到另外一個。
View本質(zhì)上是提供了對同一片存儲數(shù)據(jù)的不同看待方式,針對每種看待方式可以提供一個用戶更友好的命名(例如a, c),但實際并不影響底層數(shù)據(jù)的存儲。對任何張量,在PyTorch中永遠(yuǎn)存在一對Tensor-Storage的概念。其中Tensor描述的就是張量,包含邏輯大小、步幅以及offset,而Storage則描述的是物理存儲,包含了物理存儲大小、單元數(shù)據(jù)類型(dtype)等。
Tensor Operations
在最抽象的層面上,任何的張量操作,都會涉及到兩次調(diào)度(dispatches):
第一次調(diào)度源于設(shè)備類型(device type)和張量布局(tensor layout)。對于不同設(shè)備存儲類型(CPU or CUDA)以及不同布局類型(strided or sparse)的張量,其操作(加減乘除等)的底層實現(xiàn)方式完全不同,并且底層計算分別位于不同的動態(tài)庫中(e.g. libcaffe2.so or libcaffe2_gpu.so)。因此第一次調(diào)度是動態(tài)調(diào)度,需要通過虛函數(shù)方式調(diào)用。
第二次調(diào)度源于元素的數(shù)據(jù)類型。即使在同一個設(shè)備類型上,不同類型數(shù)據(jù)(float or int)的實現(xiàn)方式也不相同。通過簡單的switch-statement的方式實現(xiàn)不同類型的支持。
Tensor Extensions
PyTorch中的張量由如下三位一體的獨立屬性描述:
設(shè)別類型:描述了張量實際的物理存儲位置,例如GPU、NVIDIA GPU(cuda)、AMD GPU(hip),或者TPU(xla)。之所以將設(shè)備類型作為一個顯著的特性,是因為不同設(shè)備上的內(nèi)存分配方式可能完全不同;
布局類型:描述了如何從邏輯上解釋物理存儲方式,包括了最常用的strided布局、以及其它類型,比如稀疏布局(需要通過一對索引-數(shù)據(jù)描述),以及更為復(fù)雜的MKL-DNN張量布局等。
數(shù)據(jù)類型:描述了數(shù)據(jù)單元的實際類型,可以是float、int、以及量化整數(shù)等。
對于任何需要自定義的張量,均可以通過擴展如上三個獨立屬性來實現(xiàn)。另外一種更為便捷的方法是通過wrapper的方式封裝PyTorch已經(jīng)定義的張量,實現(xiàn)自定義張量。采用何種方式取決于自定義的張量是否會參與自動梯度的反向計算。例如,對于稀疏張量顯然是需要通過獨立擴展實現(xiàn),因為在任何涉及embedding的網(wǎng)絡(luò)中,embedding的梯度也需要是稀疏的。
為了提高屬性訪問效率,所有類型張量都具備的屬性按照固定的布局定義在張量結(jié)構(gòu)體的起始位置,包括size、dtype、type_id等。此外與最常用的strided布局的張量相關(guān)的屬性,包括strides、storage等也包含在固定的布局中,而其它與特定類型張量相關(guān)的屬性則定義在自定義域中。核心屬性的固定布局設(shè)計避免了使用類似于虛函數(shù)訪問的開銷。