mxnet中conv操作的實施

之前打算留一段時間把conv操作理解下，但似乎V0.10與之前的版本有些改動，實現起來變容易了些（可能是這樣的吧）。
這里關注的是前向過程。

Code 與 Note

二級標題按照調用路線，再次的標題按照Code和Note進行，另外的一些Note標在Code中。

Forward

Code

//@src/operator/convolution-inl.h
  virtual void Forward(const OpContext &ctx,
                       const std::vector<TBlob> &in_data,
                       const std::vector<OpReqType> &req,
                       const std::vector<TBlob> &out_data,
                       const std::vector<TBlob> &aux_args) {
    using namespace mshadow;
    using namespace mshadow::expr;
    CHECK_EQ(req[conv::kOut], kWriteTo);
    size_t expected = param_.no_bias ? 2 : 3;
    CHECK_EQ(in_data.size(), expected);
    CHECK_EQ(out_data.size(), 1U);
    CHECK_EQ(req[conv::kOut], kWriteTo);
    LayerSetUp(in_data[conv::kData].shape_, out_data[conv::kOut].shape_);
    Stream<xpu>* s = ctx.get_stream<xpu>();
    // allocate workspace for col_buffer
    Tensor<xpu, 1, DType> workspace = ctx.requested[conv::kTempSpace]
      .get_space_typed<xpu, 1, DType>(Shape1(col_buffer_size_), s);
    // calculate the shape of col_buffer
    TShape col_buffer_shape(num_spatial_axes_ + 1);           // num_spatial_axes = kernel.ndim =2
    col_buffer_shape[0] = conv_in_channels_ * param_.kernel.Size();   // 單個 filter 的大小
    for (index_t i = 1; i < col_buffer_shape.ndim(); ++i) {
      col_buffer_shape[i] = out_data[0].shape_[i+1];                               //  col_buffer_shape[-2:] = [out_H,out_W]
    }
    // create a column buffer using workspace and col_buffer_shape
    TBlob col_buffer(workspace.dptr_, col_buffer_shape, xpu::kDevMask, DataType<DType>::kFlag);

    // initialize weight and col_buffer 3D tensors for using gemm

    index_t M = conv_out_channels_ / group_;                      // group_ 的default 值為1， M 即 filter 的數量
    index_t N = conv_out_spatial_dim_;                                 //  H x W
    index_t K = kernel_dim_;                                                  // conv_in_channel x kernel.Size()

    Tensor<xpu, 3, DType> weight_3d = in_data[conv::kWeight].get_with_shape<xpu, 3, DType>(
      Shape3(group_, M, K), s);
    Tensor<xpu, 3, DType> col_buffer_3d = col_buffer.get_with_shape<xpu, 3, DType>(
      Shape3(group_, K, N), s);
    Tensor<xpu, 4, DType> output_4d = out_data[conv::kOut].get_with_shape<xpu, 4, DType>(
      Shape4(num_, group_, M, N), s);
    for (index_t n = 0; n < num_; ++n) {                                //   num_ = batch_size
      // transform image to col_buffer in order to use gemm
      im2col(s, in_data[conv::kData].dptr<DType>()+n*input_dim_, in_data[conv::kData].shape_,
             col_buffer.shape_, param_.kernel, param_.pad, param_.stride, param_.dilate,
             col_buffer.dptr<DType>());
      Tensor<xpu, 3, DType> output_3d = output_4d[n];
      for (index_t g = 0; g < group_; ++g) {
        ASSIGN_DISPATCH(output_3d[g], req[conv::kOut], dot(weight_3d[g], col_buffer_3d[g]));
      }
    }
    if (bias_term_) {
      Tensor<xpu, 1, DType> bias = in_data[conv::kBias].get<xpu, 1, DType>(s);
      Tensor<xpu, 3, DType> output_3d = out_data[conv::kOut].get_with_shape<xpu, 3, DType>(
        Shape3(num_, conv_out_channels_, conv_out_spatial_dim_), s);
      // has bias term, broadcast it to the same shape of output_3d in channel dim
      output_3d += mshadow::expr::broadcast<1>(bias, output_3d.shape_);
    }
  }

Note

內積方式是其思想。
index_t K暗示的實現思路是，將一個filter的weight全部提出，然后再將輸入data中的對應元素提出，兩者進行內積，有多少個filter就進行多少次內積（注意filter的個數正是輸出map的channel數，index_t M）
index_t N說明，計算是以輸出的map中每個pixel為基本單位進行。
之所以這里特地把這個思路關注一下，是因為之前自己實現的一個conv操作是通過稀疏矩陣與向量做內積進行的(Julia中)，這樣實現的一個顯著缺點是內存消耗增長過快；即使是稀疏矩陣，要存放以pixel數量的值為長度的方陣，仍然會將通用性達到谷底，16G的內存，單個的\(15\times15\)左右的核已經使內存接近崩潰了。
做個對比，用稀疏矩陣的方法的一個主要的缺陷是重復的數據占用了大量的內存（\(H\times W\)列中占據內存的全是一個核的參數）。
后面來討論這個思路的情況如何。

重點是兩個數據。
weight_3d的尺寸:：\(1\times M\times K\);
col_buffer_3d的尺寸：\(1\times K\times N\)
以上默認group_=1。
我們再來回顧一下M,K,N的意義。
M代表了filter的數量（也是輸出數據，channel的數量）；
K是單個filter的參數數量；
N是輸出中，單個channel中二維feature map的大小。
weight_3d中存放的，顯然是所有filter的weight；而col_buffer_3d，按照這種思路，應該用來存放對應的輸入元素。直覺上，觀察一下兩個數據的尺寸（在二維角度下），可以發現其內積的尺寸就是\(M\times N\)，這正好就是一個樣本對應的輸出尺寸。
於是im2col便是用來將輸入數據排列成對應的順序，放到col_buffer_3d中。

im2col

這里關注GPU上的實現，在src/operator/nn/im2col.cuh中

Code

//@src/operator/nn/im2col.cuh

template <typename DType>
inline void im2col(mshadow::Stream<gpu>* s,
                   const DType* data_im, const TShape& im_shape,
                   const TShape& col_shape, const TShape& kernel_shape,
                   const TShape& pad, const TShape& stride,
                   const TShape& dilation, DType* data_col) {
  // num_axes should be smaller than block size
  index_t num_spatial_axes = kernel_shape.ndim();
  CHECK_LT(num_spatial_axes, mshadow::cuda::kBaseThreadNum);
  index_t num_kernels = im_shape[1] * col_shape.ProdShape(1, col_shape.ndim());   

Note

這是前半部分，這里關注下index_t num_kernels。此處的num_kernel是指希望啟動的CUDA中的kernel數量（實際分配數量受硬件限制，詳見之后的cuda_get_num_blocks(num_kernels)）。
但觀察下其值，可以猜測啟動的每個kernel要做的事情。它的值是\(conv_in_channel\times out_H\times out_W=K\times N/k.Size()\)，而這個函數的目的是要完成對col_buffer_3d的填充，所以每個線程kernel將對一個卷積kernel所覆蓋的in_data中的數據進行排列操作。

im2col_gpu_kernel

im2col后半部分就是啟動這個函數，所以直接到這里。這個函數就是之前猜測的實現，就不進行單獨Note了，Code里面的注釋可以提示進行操作。

//@src/operator/nn/im2col.cuh
template <typename DType>
__global__ void im2col_gpu_kernel(const int n, const DType* data_im,
    const int height, const int width, const int kernel_h, const int kernel_w,
    const int pad_h, const int pad_w,
    const int stride_h, const int stride_w,
    const int dilation_h, const int dilation_w,
    const int height_col, const int width_col,
    DType* data_col) {
/*
height_col= out_data.H
width_col = ~~      .W

n         = in_data.channel x out_data.H x out_data.W

data_col  : [ group , in_data.channel x kernel.size/group   , out_data.H x out_data.W ]
            [   1   ,                  K                    ,            N            ]

so, need  x kernel.size in this funcn...
*/
  CUDA_KERNEL_LOOP(index, n) {
    const int h_index = index / width_col;           //  in_data.ch_th x out_data.H_th
    const int h_col = h_index % height_col;          // out_data.H_th
    const int w_col = index % width_col;             // out_data.W_th
    const int c_im = h_index / height_col;           // in_data.ch_th
    const int c_col = c_im * kernel_h * kernel_w;    // see data_col's layout above, weight_3d : [ conv_out_ch, K ] so, ... dot(weight_3d, data_col)
    const int h_offset = h_col * stride_h - pad_h;   // (h_offset, w_offset) is the top-left point, where facing its counterpart in kernel
    const int w_offset = w_col * stride_w - pad_w;
    DType* data_col_ptr = data_col;
    data_col_ptr += (c_col * height_col + h_col) * width_col + w_col;  // see data_col's layout above
    const DType* data_im_ptr = data_im;
    data_im_ptr += (c_im * height + h_offset) * width + w_offset;    // see data_col's layout above
    for (int i = 0; i < kernel_h; ++i) {
      for (int j = 0; j < kernel_w; ++j) {
        int h_im = h_offset + i * dilation_h;
        int w_im = w_offset + j * dilation_w;
        *data_col_ptr =
            (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) ?
            data_im_ptr[i * dilation_h * width + j * dilation_w] : static_cast<DType>(0);
        data_col_ptr += height_col * width_col;       //    +N  (next row)
      }
    }
  }
}

小結

1. conv的實現是通過內積進行的，而gpu在此處的加速作用，是並行地搬運數據（然而由於存在重復的數據，沒有利用局部存儲將會損失一部分性能）和內積操作。
2. 回到之前的那個話題，用稀疏矩陣實現的conv操作，要對kernel size大小的數據進行img size的重復，這導致內存消耗攀升顯著。但正在討論中的這個方法也存在着數據重復問題：每個in_data中的數據將被重復kernel size數量次（因為filter在產生每個輸出pixel中，都按照自身的需要排列輸入數據，每個輸入中的數據將在kernel size的范圍被考慮。而這種排列是在計算前就已經完成了的——將為重復的數據付出內存代價）。正式地來看，重復數據量分別為：\(k\times k\times H\times W\)，\(H\times W\times k\times k\)。也就是說兩者沒有差別。
3. 從程序上來看，im2col啟動了線程kernel，但(我似乎記得)CUDA的kernel調用只會顯式地阻塞，這就為mxnet的非阻塞式設計提供了便利。

跟進

1. MSHADOW_CUDA_POST_KERNEL_CHECK 的使用。