不是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层面也有一个对应的函数:
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的数据结构,原来矩阵库都是一毛一样的。。。