mshadow入门指南

littleorange 2020-11-01 原文

这篇文章简单介绍了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申请的GPU显存

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的高性能计算能力。