mshadow入门指南

mshadow是一个基于表达式模板实现的张量库，在MXNet框架中被广泛使用。这篇文章简单介绍了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的高性能计算能力。

mshadow还提供了一个名为TensorContainer的STL风格容器，它的行为和张量类似，但是会在析构时自动释放内存。

逐元素的操作

mshadow中所有的运算符（+，-，*，/等）都是元素级操作。考虑如下SGD更新代码：

void UpdateSGD(Tensor<cpu, 2> weight, Tensor<cpu, 2> grad, float eta, float lambda) {
  weight -= eta * (grad + lambda * weight);
}

在编译期，上述代码会被转化成下面的代码：

void UpdateSGD(Tensor<cpu,2> weight, Tensor<cpu,2> grad, float eta, float lambda) {
  for (index_t y = 0; y < weight.size(0); ++y) {
    for (index_t x = 0; x < weight.size(1); ++x) {
      weight[y][x] -= eta * (grad[y][x] + lambda * weight[y][x]);
    }
  }
}

可以看到，代码转换过程中没有发生任何内存分配操作。对于Tensor<gpu, k>，对应的函数会被转化成具有相同含义的CUDA核函数。使用表达式模板，上述的转换过程会发生在编译期。

CPU与GPU上的通用代码

由于mshadow对Tensor <cpu，k>和Tensor <gpu，k>提供了相同的接口，因此我们可以轻松地编写运行在CPU和GPU上地代码。比如，下面的代码可以同时被CPU和GPU上的张量所使用。

template<typename xpu>
void UpdateSGD(Tensor<xpu, 2> weight, const Tensor<xpu, 2> &grad,
               float eta, float lambda) {
  weight -= eta * (grad + lambda * weight);
}

矩阵乘法

mshadow提供了一个矩阵点积的实现，内部封装了MKL和cuBLAS等库。

template<typename xpu>
void Backprop(Tensor<xpu, 2> gradin,
              const Tensor<xpu, 2> &gradout,
              const Tensor<xpu, 2> &netweight) {
  gradin = dot(gradout, netweight.T());
}

用户自定义操作

假设用户要在mshadow中自定义一个逐元素的sigmoid函数，那么我们可以通过下面的代码将sigmoid操作加入到mshadow中。

struct sigmoid {
  MSHADOW_XINLINE static float Map(float a) {
    return 1.0f / (1.0f + expf(-a));
  }
};
template<typename xpu>
void ExampleSigmoid(Tensor<xpu, 2> out, const Tensor<xpu, 2> &in) {
  out = F<sigmoid>(in * 2.0f) + 1.0f;
}

转换后的代码就如下所示（CPU版本）：

template<typename xpu>
void ExampleSigmoid(Tensor<xpu, 2> out, const Tensor<xpu, 2> &in) {
  for (index_t y = 0; y < out.size(0); ++y) {
    for(index_t x = 0; x < out.size(1); ++x) {
      out[y][x] = sigmoid::Map(in[y][x] * 2.0f) + 1.0f;
    }
  }
}

同样，我们也可以定义形如out = F<sigmoid>+2.0以及out = F<sigmoid>(F<sigmoid>(in))的复合表达式。此外，在GPU上运行的版本将会被转化为CUDA和函数，详见defop.cpp文件。

完整的例子

下面的代码来自于basic.cpp，它展示了如何使用mshadow进行计算。

// header file to use mshadow
#include "mshadow/tensor.h"
// this namespace contains all data structures, functions
using namespace mshadow;
// this namespace contains all operator overloads
using namespace mshadow::expr;

int main(void) {
  // intialize tensor engine before using tensor operation, needed for CuBLAS
  InitTensorEngine<cpu>();
  // assume we have a float space
  float data[20];
  // create a 2 x 5 x 2 tensor, from existing space
  Tensor<cpu, 3> ts(data, Shape3(2,5,2));
    // take first subscript of the tensor
  Tensor<cpu, 2> mat = ts[0];
  // Tensor object is only a handle, assignment means they have same data content
  // we can specify content type of a Tensor, if not specified, it is float bydefault
  Tensor<cpu, 2, float> mat2 = mat;

  // shape of matrix, note size order is the same as numpy
  printf("%u X %u matrix\n", mat.size(0), mat.size(1));

  // initialize all element to zero
  mat = 0.0f;
  // assign some values
  mat[0][1] = 1.0f; mat[1][0] = 2.0f;
  // elementwise operations
  mat += (mat + 10.0f) / 10.0f + 2.0f;

  // print out matrix, note: mat2 and mat1 are handles(pointers)
  for (index_t i = 0; i < mat.size(0); ++i) {
    for (index_t j = 0; j < mat.size(1); ++j) {
      printf("%.2f ", mat2[i][j]);
    }
    printf("\n");
  }
  // shutdown tensor enigne after usage
  ShutdownTensorEngine<cpu>();
  return 0;
}