pytorch中的底層很多代碼都是來(lái)源于的torch的低層Tensor庫(kù)
TH = TorcH
THC = TorcH Cuda
THCS = TorcH Cuda Sparse
THCUNN = TorcH CUda Neural Network (see cunn)
THD = TorcH Distributed
THNN = TorcH Neural Network
THS = TorcH Sparse
特別推薦查看Torch7的官方介紹Tensor
其他參考：pytorch作者adam paszke的博客

前面我們介紹strided indexing scheme時(shí)就已經(jīng)說(shuō)過(guò)，很多矩陣(ndarray)庫(kù)都是存儲(chǔ)和表示分開(kāi)，以便很多變量共享內(nèi)存，TH庫(kù)也是這樣

THTensor, THStorage, THAllocator的關(guān)系##

先整體概覽一下他們3者的關(guān)系:

數(shù)據(jù)存儲(chǔ): THStorage

所有在CPU上的張量實(shí)際上都是內(nèi)存中的一個(gè)一維C數(shù)組（C指針）data來(lái)存儲(chǔ)，并且使用引用計(jì)數(shù)（reference count）來(lái)管理內(nèi)存。

typedef struct THStorage {real *data;ptrdiff_t size;int refcount; // 引用計(jì)數(shù)char flag;THAllocator *allocator;void *allocatorContext; // GPU或CPU的上下文，在GPU時(shí)有用，在CPU暫無(wú)用到struct THStorage *view; } THStorage; // flag是4位標(biāo)識(shí)符號(hào),例如flag=0b0101 代表可以釋放內(nèi)存且采用了引用計(jì)數(shù) ， // 所以在storage釋放內(nèi)存時(shí)首先判斷是否可以釋放，然后如果采用了引用計(jì)數(shù)，只有計(jì)數(shù)為0時(shí)才可以釋放內(nèi)存。 #define TH_STORAGE_REFCOUNTED 1 #define TH_STORAGE_RESIZABLE 2 #define TH_STORAGE_FREEMEM 4 #define TH_STORAGE_VIEW 8

所有構(gòu)造新THStorage的函數(shù)都以new開(kāi)頭，后面跟具有相關(guān)含義的后綴名。

// 空的THStorage -> return THStorage_(newWithSize)(0); TH_API THStorage* THStorage_(new)(void); // 指定大小的THStorage -> return THStorage_(newWithAllocator)(size, &THDefaultAllocator, NULL); // 未賦值，直接malloc了內(nèi)存，未初始化這塊內(nèi)存 TH_API THStorage* THStorage_(newWithSize)(ptrdiff_t size); TH_API THStorage* THStorage_(newWithSize1)(real); TH_API THStorage* THStorage_(newWithSize2)(real, real); TH_API THStorage* THStorage_(newWithSize3)(real, real, real); TH_API THStorage* THStorage_(newWithSize4)(real, real, real, real); TH_API THStorage* THStorage_(newWithMapping)(const char *filename, ptrdiff_t size, int flags);/* takes ownership of data */ // 生成一個(gè)THStorage，其data直接指向傳入的data，不再重新開(kāi)辟內(nèi)存 // return THStorage_(newWithDataAndAllocator)(data, size, &THDefaultAllocator, NULL); TH_API THStorage* THStorage_(newWithData)(real *data, ptrdiff_t size);TH_API THStorage* THStorage_(newWithAllocator)(ptrdiff_t size,THAllocator* allocator,void *allocatorContext); TH_API THStorage* THStorage_(newWithDataAndAllocator)(real* data, ptrdiff_t size, THAllocator* allocator, void *allocatorContext);

內(nèi)存分配: THAllocator

// THAllocator.h /** THAllocator 定義了3個(gè)函數(shù)指針 malloc realloc free* 即指向函數(shù)的指針,在THAllocator.c中會(huì)把具體的函數(shù)地址賦給它們*/ typedef struct THAllocator {void* (*malloc)(void*, ptrdiff_t);void* (*realloc)(void*, void*, ptrdiff_t);void (*free)(void*, void*); } THAllocator;// THAllocator.c // 它最終實(shí)際調(diào)用的就是#include <malloc.h>中的malloc(size) // 但是它這么套了幾層，是為了考慮各個(gè)平臺(tái)上malloc時(shí)數(shù)據(jù)對(duì)齊的問(wèn)題，這里不再深究 static void *THDefaultAllocator_alloc(void* ctx, ptrdiff_t size) {return THAlloc(size); } // realloc(ptr, size) static void *THDefaultAllocator_realloc(void* ctx, void* ptr, ptrdiff_t size) {return THRealloc(ptr, size); } // free(ptr) static void THDefaultAllocator_free(void* ctx, void* ptr) {THFree(ptr); } THAllocator THDefaultAllocator = {&THDefaultAllocator_alloc,&THDefaultAllocator_realloc,&THDefaultAllocator_free };

數(shù)據(jù)查看(表示): THTensor

typedef struct THTensor {// 一個(gè)用來(lái)存儲(chǔ)張量各個(gè)維度大小的一維數(shù)組int64_t *size;// 一個(gè)用來(lái)存儲(chǔ)張量各個(gè)角標(biāo)偏移量的數(shù)組，之前strided indexing scheme介紹過(guò)int64_t *stride;// 維度int nDimension; // Note: storage->size may be greater than the recorded size// of a tensor// 持有的實(shí)際存儲(chǔ)的指針，存儲(chǔ)大小可能實(shí)際上是大于于等于張量大小THStorage *storage; // 存儲(chǔ)偏移ptrdiff_t storageOffset;// 引用計(jì)數(shù)int refcount;// 同THStoragechar flag;} THTensor;

如果通過(guò)THTensor去查看數(shù)據(jù)實(shí)際存儲(chǔ)地THStorage呢？我們?cè)趕trided indexing scheme這章里講得很清楚了，主要就通過(guò)storageOffset和stride。

這個(gè)圖片還缺了一個(gè)storageOffset未加上去
我們以torch.narrow這個(gè)函數(shù)為例，解讀如何通過(guò)storageOffset和stride實(shí)現(xiàn)THStorage的不同視圖THTensor

# narrow(dimension, start, length) → Tensor # 抽取Tensor某個(gè)維度上的一部分 # stride: 3,1 size: 3,3 x = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])1 2 34 5 67 8 9 # 抽取x第0個(gè)維度，從0開(kāi)始，長(zhǎng)度為2，所有就是第0行和第1行 # dimension=0 start=0 length=2 # self->storageOffset += firstIndex*self->stride[dimension]; # self->size[dimension] = size; # 可以看到storageOffset=0*3=0，不變 # size: length,3 = 2,3 # stride不變 x.narrow(0, 0, 2)1 2 34 5 6 [torch.FloatTensor of size (2,3)] # 抽取x第1個(gè)維度，從1開(kāi)始，長(zhǎng)度為2，所有就是第1行和第2行 # dimension=1 start=1 length=2 # self->storageOffset += firstIndex*self->stride[dimension]; # self->size[dimension] = size; # 可以看到storageOffset=1*1=1,其起始偏移變?yōu)榱?，剛好把第一個(gè)元素1跳過(guò)去 # size: 3,length = 3,2 # stride不變 x.narrow(1, 1, 2)2 35 68 9 [torch.FloatTensor of size (3,2)]

低層實(shí)現(xiàn)如下:

void THTensor_(narrow)(THTensor *self, THTensor *src, int dimension, int64_t firstIndex, int64_t size) {if(!src)src = self;THArgCheck( (dimension >= 0) && (dimension < src->nDimension), 2, "out of range");THArgCheck( (firstIndex >= 0) && (firstIndex < src->size[dimension]), 3, "out of range");THArgCheck( (size > 0) && (firstIndex <= src->size[dimension] - size), 4, "out of range");THTensor_(set)(self, src);if(firstIndex > 0)self->storageOffset += firstIndex*self->stride[dimension];self->size[dimension] = size; }

---

PyTorch源碼淺析（一）
PyTorch – Internal Architecture Tour
機(jī)器之心翻譯: PyTorch – Internal Architecture Tour
torch.Tensor

總結(jié)

以上是生活随笔為你收集整理的TH库学习(一): THTensor, THStorage, THAllocator介绍的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。