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

本站文章如无特殊说明，均为本站原创，如若转载，请注明出处：对pytorch中Tensor的剖析 - Python技术站

对pytorch中Tensor的剖析

相关文章