这篇文章简单介绍了mshadow的基本用法和特性，文章主要翻译自mshadow/guide/README。
张量数据结构
mshadow中的主要数据结构就是张量（Tensor），下面是一个简化版本的声明定义（来自mshadow/tensor.h文件）：
typedef unsigned index_t;
template<int dimension>
struct Shape {
index_t shape_[dimension];
};
template<typename Device, int dimension, typename DType = float>
struct Tensor {
DType *dptr_;
Shape<dimension> shape_;
Stream<Device> stream_;
index_t stride_;
};
// this is how shape object declaration look like
Shape<2> shape2;
// this is how tensor object declaration look like
Tensor<cpu, 2> ts2;
Tensor<gpu, 3, float> ts3;
在上述代码中，
Tensor<cpu,2>
是内存上的一个二维张量，而
Tensor<gpu,3>
是存储在GPU显存上的一个三维张量。
Shape<k>
给出了一个k维张量的维度信息。通过使用模板编程技术，用户可以申请存储在特定设备上的不同尺寸的张量。下面是一个二维张量的定义：
struct Shape<2> {
index_t shape_[2];
};
struct Tensor<cpu, 2, float> {
float *dptr_;
Shape<2> shape_;
index_t stride_;
};
-
Tensor<cpu, 2>
dptr_
-
Shape<2>
-
stride_
stride_
下面的代码可以帮助我们更好地理解mahsdow中的张量。
float data[9] = {0, 1, 2, 3, 4, 5, 6, 7, 8};
Tensor<cpu, 2> ts;
ts.dptr_ = data;
ts.shape_ = mshadow::Shape2(3, 2);
ts.stride_ = 3;
// now: ts[0][0] == 0, ts[0][1] == 1 , ts[1][0] == 3, ts[1][1] == 4
for (index_t i = 0; i < ts.size(0); ++i) {
for (index_t j = 0; j < ts.size(1); ++j) {
printf("ts[%u][%u]=%f\n", i, j, ts[i][j]);
}
}
代码中的
ts
是一个\(3 \times 2\)的矩阵，其中
data[2]
、
data[5]
以及
data[8]
作为填充单元被忽略掉。如果想访问连续内存，设置
stride_=shape_[1]
即可。
内存分配
mshadow的一个重要设计就是将张量视作一个“白盒”。只要我们把
dptr_
和
shape_
以及
stride_
对应起来，它就可以工作：
- 对于
Tensor<cpu, k>
dptr_
new float[]
- 对于
Tensor<gpu, k>
dptr_
cudaMallocPitch
mshadow提供了显式内存分配的函数，如下所示：
// create a 5 x 3 tensor on the device, and allocate space
Tensor<gpu, 2> ts2(Shape2(5, 3));
AllocSpace(&ts2);
// allocate 5 x 3 x 2 tensor on the host, initialized by 0
Tensor<cpu, 3> ts3 = NewTensor<cpu>(Shape3(5,3,2), 0.0f);
// free space
FreeSpace(&ts2); FreeSpace(&ts3);
mshadow中的所有的内存分配操作都是显式进行的，不会出现任何隐式的内存分配或内存销毁等操作。这就意味着，
Tensor<cpu, k>
更像一个指针（或引用），而不是一个对象。如果我们把一个张量赋值给另一个，那么他们会指向相同的内存空间。另外，这种特性对用户来说是十分友好的，只需要把一个指针交给mshadow，即可零成本地受益于mshadow的高性能计算能力。