对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中

 typedef struct THStorage
 {
     real *data;
     ptrdiff_t size;
     int refcount;
     char flag;
     THAllocator *allocator;
     void *allocatorContext;
     struct THStorage *view;
 } 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中

 typedef struct THTensor
 {
     int64_t *size;   // 注意是指针
     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;
     char flag;
 } 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:
Tensor b:       x  x      x  x      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 中。

 int THTensor_(isContiguous)(const THTensor *self)
 {
   int64_t z = ;
   int d;
   for(d = self->nDimension-; d >= ; d--)
   {
     if(self->size[d] != )
     {
       if(self->stride[d] == z)
         z *= self->size[d]; // 如果是连续的，应该在这循环完然后跳到下面return 1
       else
         return ;
     }
   }
   return ;
 }

把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的数据结构，原来矩阵库都是一毛一样的。。。