pytorch中的底層很多代碼都是來源于的torch的低層Tensor庫
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時就已經說過，很多矩陣(ndarray)庫都是存儲和表示分開，以便很多變量共享內存，TH庫也是這樣

THTensor, THStorage, THAllocator的關系##

先整體概覽一下他們3者的關系:

數據存儲: THStorage

所有在CPU上的張量實際上都是內存中的一個一維C數組（C指針）data來存儲，并且使用引用計數（reference count）來管理內存。

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

所有構造新THStorage的函數都以new開頭，后面跟具有相關含義的后綴名。

// 空的THStorage -> return THStorage_(newWithSize)(0); TH_API THStorage* THStorage_(new)(void); // 指定大小的THStorage -> return THStorage_(newWithAllocator)(size, &THDefaultAllocator, NULL); // 未賦值，直接malloc了內存，未初始化這塊內存 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 */ // 生成一個THStorage，其data直接指向傳入的data，不再重新開辟內存 // 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);

內存分配: THAllocator

// THAllocator.h /** THAllocator 定義了3個函數指針 malloc realloc free* 即指向函數的指針,在THAllocator.c中會把具體的函數地址賦給它們*/ typedef struct THAllocator {void* (*malloc)(void*, ptrdiff_t);void* (*realloc)(void*, void*, ptrdiff_t);void (*free)(void*, void*); } THAllocator;// THAllocator.c // 它最終實際調用的就是#include <malloc.h>中的malloc(size) // 但是它這么套了幾層，是為了考慮各個平臺上malloc時數據對齊的問題，這里不再深究 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 };

數據查看(表示): THTensor

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

如果通過THTensor去查看數據實際存儲地THStorage呢？我們在strided indexing scheme這章里講得很清楚了，主要就通過storageOffset和stride。

這個圖片還缺了一個storageOffset未加上去
我們以torch.narrow這個函數為例，解讀如何通過storageOffset和stride實現THStorage的不同視圖THTensor

# narrow(dimension, start, length) → Tensor # 抽取Tensor某個維度上的一部分 # 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個維度，從0開始，長度為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個維度，從1開始，長度為2，所有就是第1行和第2行 # dimension=1 start=1 length=2 # self->storageOffset += firstIndex*self->stride[dimension]; # self->size[dimension] = size; # 可以看到storageOffset=1*1=1,其起始偏移變為了1，剛好把第一個元素1跳過去 # size: 3,length = 3,2 # stride不變 x.narrow(1, 1, 2)2 35 68 9 [torch.FloatTensor of size (3,2)]

低層實現如下:

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
機器之心翻譯: PyTorch – Internal Architecture Tour
torch.Tensor

總結

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

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。