对pytorch中Tensor的剖析

不是python层面Tensor的剖析,是C层面的剖析。

 

看pytorch下lib库中的TH好一阵子了,TH也是torch7下面的一个重要的库。

可以在torch的github上看到相关文档。看了半天才发现pytorch借鉴了很多torch7的东西。

pytorch大量借鉴了torch7下面lua写的东西并且做了更好的设计和优化。

https://github.com/torch/torch7/tree/master/doc

 

pytorch中的Tensor是在TH中实现的。TH = torch

TH中先实现了一个THStorage,再在THStorage的基础上实现了THTensor。

THStorage定义如下,定义在TH/generic/THStorage.h中

 1 typedef struct THStorage
 2 {
 3     real *data;
 4     ptrdiff_t size;
 5     int refcount;
 6     char flag;
 7     THAllocator *allocator;
 8     void *allocatorContext;
 9     struct THStorage *view;
10 } THStorage;

这些成员里重点关注*data和size就可以了。

real *data中的real会在预编译的时候替换成预先设计的数据类型,比如int,float,byte等。

比如 int a[3] = {1,2,3},data是数组a的地址,对应的size是3,不是sizeof(a)。

所以*data指向的是一段连续内存。是一维的!

 

讲Tensor前先回顾下数组在内存中的排列方式。参看《C和指针》8.2节相关内容。

比如 int a[3][6]; 内存中的存储顺序为:

00 01 02 03 04 05 10 11 12 13 14 15 20 21 22 23 24 25 

是连续存储的。存储顺序按照最右边的下标率先变化。

 

然后数组a是2维的,nDimension = 2。dimension从0开始算起。

 

size(a) = {3,6}
[3] 是 dimension 0    size[0] = 3
[6] 是 dimension 1    size[1] = 6
nDimension = 2

 

THTensor定义如下,定义在TH/generic/THTensor.h中

 

 1 typedef struct THTensor
 2 {
 3     int64_t *size;   // 注意是指针
 4     int64_t *stride; // 注意是指针
 5     int nDimension;
 6 
 7     // Note: storage->size may be greater than the recorded size
 8     // of a tensor
 9     THStorage *storage;
10     ptrdiff_t storageOffset;
11     int refcount;
12     char flag;
13 } THTensor;

 

比如

z = torch.Tensor(2,3,4)   // 新建一个张量,size为 2,3,4

size(z) = {2,3,4}
[2] 是 dimension 0    size[0] = 2
[3] 是 dimension 1    size[1] = 3
[4] 是 dimension 2    size[2] = 4
nDimension = 3

THStorage只管理内存,是一维的。

THTensor通过size和nDimension将THStorage管理的一维内存映射成逻辑上的多维张量,

底层还是一维的。但是注意,代表某个Tensor的底层内存是一维的但是未必是连续的!

 

把Tensor按照数组来理解好了。

Tensor a[3][6]  裁剪(narrow函数)得到一个 Tensor b[3][4],在内存中就是

Tensor a:  00 01 02 03 04 05 10 11 12 13 14 15 20 21 22 23 24 25 
Tensor b:  00 01 02 03  x  x 10 11 12 13  x  x 20 21 22 23  x  x

narrow函数并不会真正创建一个新的Tensor,Tensor b还是指向Tensor a的那段内存。

所以Tensor b在内存上就不是连续的了。

 

那么怎么体现Tensor在内存中是连续的呢?就靠THTensor结构体中的

size,stride,nDimension共同判断了。

pytorch的Tensor有个 contiguous 函数,C层面也有一个对应的函数:

int THTensor_(isContiguous)(const THTensor *self)
判断 Tensor 在内存中是否连续。定义在 TH/generic/THTensor.c 中。
 1 int THTensor_(isContiguous)(const THTensor *self)
 2 {
 3   int64_t z = 1;
 4   int d;
 5   for(d = self->nDimension-1; d >= 0; d--)
 6   {
 7     if(self->size[d] != 1)
 8     {
 9       if(self->stride[d] == z)
10         z *= self->size[d]; // 如果是连续的,应该在这循环完然后跳到下面return 1
11       else
12         return 0;
13     }
14   }
15   return 1;
16 }

把Tensor a[3][6] 作为这个函数的参数:

size[0] = 3    size[1] = 6    nDimension = 2      z =1
d = 1   if size(1) = 6 != 1   if stride[1] == 1   z = z*size(d)=6
d = 0   if size(0) = 3 != 1   if stride[0] == 6   z = z*size(d)=6*3 = 18
因此,对于连续存储的a
stride = {6,1}
size = {3,6}

 

再举一个Tensor c[2][3][4]的例子,如果c是连续存储的,则:

stride = {12,4,1}
size =    { 2,3,4}  // 2所对应的stride就是 右边的数相乘(3x4), 3所对应的stride就是右边的数相乘(4)

stride(i)返回第i维的长度。stride又被翻译成步长。

比如第0维,就是[2]所在的维度,Tensor c[ i ][ j ][ k ]跟Tensor c[ i+1 ][ j ][ k ]

在连续内存上就距离12个元素的距离。

 

对于内存连续的stride,计算方式就是相应的size数右边的数相乘。

 

所以不连续呢?

对于a[3][6]

stride = {6,1} 
size =   {3,6}

对于从a中裁剪出来的b[3][4]

stride = {6,1} 
size =   {3,4}

stride和size符合不了 右边的数相乘 的计算方法,所以就不连续了。

 

所以一段连续的一维内存,可以根据size和stride 解释 成  逻辑上变化万千,内存上是否连续 的张量。

比如24个元素,可以解释成 4 x 6 的2维张量,也可以解释成 2 x 3 x 4 的3维张量。

THTensor中的 storageOffset 就是说要从 THStorage 的第几个元素开始 解释 了。

连续的内存能给程序并行化和最优化算法提供很大的便利。

其实写这篇博客是为了给理解 TH 中的 TH_TENSOR_APPLY2 等宏打基础。

这个宏就像是在C中实现了broadcast。

 

2017年12月11日01:00:22

最近意识到,用 H x W x C 和 C x H x W 哪个来装图像更好,取决于矩阵在内存中是行存储还是

列存储,这个会影响内存读取速度,进而影响算法用时。

 

后来意识到,这就是个cache-friendly的问题,大部分对程序性能的要求还上升不到要研究算法复杂度

这个地步,常规优化的话注意下缓存友好等问题就好了,再优化就要靠更专业团队写的库或者榨干硬件了。

 

看了下numpy的文档,怪不得说pytorch是numpy的gpu版本。。。

后来又看了下opencv的mat的数据结构,原来矩阵库都是一毛一样的。。。

posted on 2017-12-11 01:00  菠菜僵尸  阅读(4167)  评论(0编辑  收藏  举报