Memo
0%

给 MShadow 添加新的 Extension

发表于 分类于 ，

MShadow 是一个轻量级的 Tensor Template Library。mshadow 同时支持 C++ 和 CUDA，一套代码既可以跑在 CPU 上，也可以跑在 GPU 上。并且是 lazy computation 的。如果对于 cuda 上面的运行没有极致要求的话，mshadow 对于 Tensor 的计算是个很好的选择，性能也不会很差。如果对于性能有极致的需求，那就要手动去 tune 了，或者也可以尝试一下 tvm。
本文尝试给 mshadow 添加一个新的 extension, mshadow::expr::resize。完整实现在 Github

resize 实现了和 OpenCV 等价的双线性插值的 resize。并且实现了 2 种不同的 padding 方法：复制边缘值和直接常数 padding。
首先定义两种 padding 的枚举，方便使用。

1
2
3
namespace resize_pad {
enum PadMode { kConstant, kEdge };
}

表达式

表达式定义了 resize 操作需要的输入和必要的变量条件。真正的计算是在计算表达式的 Plan 中进行的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
template <typename SrcExp, typename DType, int srcdim>
struct ResizeExp : public MakeTensorExp<ResizeExp<SrcExp, DType, srcdim>, SrcExp, srcdim, DType> {
  const SrcExp& src_;
  float start_y_;
  float start_x_;
  float step_y_;
  float step_x_;
  index_t src_height_;
  index_t src_width_;
  index_t out_height_;
  int pad_mode_;
  DType pad_value_;
  explicit ResizeExp(const SrcExp& src, index_t out_height, index_t out_width, int pad_mode,
                     DType pad_value)
      : src_(src), out_height_(out_height), pad_mode_(pad_mode), pad_value_(pad_value) {
    Shape<srcdim> src_shape = ShapeCheck<srcdim, SrcExp>::Check(src_);
    this->shape_ = src_shape;
    this->shape_[srcdim - 2] = out_height;
    this->shape_[srcdim - 1] = out_width;
    step_y_ = src_shape[srcdim - 2] / static_cast<float>(out_height);
    step_x_ = src_shape[srcdim - 1] / static_cast<float>(out_width);
    start_y_ = (src_shape[srcdim - 2] / static_cast<float>(out_height) - 1) / 2;
    start_x_ = (src_shape[srcdim - 1] / static_cast<float>(out_width) - 1) / 2;
    src_height_ = src_shape[srcdim - 2];
    src_width_ = src_shape[srcdim - 1];
  }
};

表达式的计算

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
template <typename SrcExp, typename DType, int etype>
inline ResizeExp<SrcExp, DType, ExpInfo<SrcExp>::kDim> resize(const Exp<SrcExp, DType, etype>& src,
                                                              index_t out_height, index_t out_width,
                                                              int pad_mode = resize_pad::kEdge,
                                                              DType pad_value = 0) {
  return ResizeExp<SrcExp, DType, ExpInfo<SrcExp>::kDim>(src.self(), out_height, out_width,
                                                         pad_mode, pad_value);
}

MSHADOW_XINLINE static bool InBound(int32_t x, index_t low, index_t high) {
  return x >= low && x <= high;
}
template <typename SrcExp, typename DType, int srcdim>
struct Plan<ResizeExp<SrcExp, DType, srcdim>, DType> {
 public:
  explicit Plan(const ResizeExp<SrcExp, DType, srcdim>& e)
      : src_(MakePlan(e.src_)),
        start_y_(e.start_y_),
        start_x_(e.start_x_),
        step_y_(e.step_y_),
        step_x_(e.step_x_),
        src_height_(e.src_height_),
        src_width_(e.src_width_),
        out_height_(e.out_height_),
        pad_mode_(e.pad_mode_),
        pad_value_(e.pad_value_) {}
  MSHADOW_XINLINE DType Eval(index_t i, index_t j) const {
    const index_t dst_w = j;
    const index_t dst_h = i % out_height_;
    const index_t c = i / out_height_;
    const float src_w = start_x_ + dst_w * step_x_;
    const float src_h = start_y_ + dst_h * step_y_;
    int32_t src_h_floor = static_cast<int32_t>(std::floor(src_h));
    int32_t src_w_floor = static_cast<int32_t>(std::floor(src_w));
    int32_t src_h_ceil = src_h_floor + 1;
    int32_t src_w_ceil = src_w_floor + 1;
    if (pad_mode_ == resize_pad::kEdge) {
      auto get_src_coord = [](int32_t x, int32_t max) {
        return mshadow::op::min::Map(mshadow::op::max::Map(x, 0), max);
      };

      src_h_floor = get_src_coord(src_h_floor, src_height_ - 1);
      src_w_floor = get_src_coord(src_w_floor, src_width_ - 1);
      src_h_ceil = get_src_coord(src_h_ceil, src_height_ - 1);
      src_w_ceil = get_src_coord(src_w_ceil, src_width_ - 1);
    }

    DType top_left_value = pad_value_, top_right_value = pad_value_, bottom_left_value = pad_value_,
          bottom_right_value = pad_value_;

    if (InBound(src_h_floor, 0, src_height_ - 1) && InBound(src_w_floor, 0, src_width_ - 1)) {
      top_left_value = src_.Eval(c * src_height_ + src_h_floor, src_w_floor);
    }
    if (InBound(src_h_floor, 0, src_height_ - 1) && InBound(src_w_ceil, 0, src_width_ - 1)) {
      top_right_value = src_.Eval(c * src_height_ + src_h_floor, src_w_ceil);
    }
    if (InBound(src_h_ceil, 0, src_height_ - 1) && InBound(src_w_floor, 0, src_width_ - 1)) {
      bottom_left_value = src_.Eval(c * src_height_ + src_h_ceil, src_w_floor);
    }
    if (InBound(src_h_ceil, 0, src_height_ - 1) && InBound(src_w_ceil, 0, src_width_ - 1)) {
      bottom_right_value = src_.Eval(c * src_height_ + src_h_ceil, src_w_ceil);
    }
    const float dy = src_h - src_h_floor;
    const float dx = src_w - src_w_floor;
    float result = top_left_value * (1 - dy) * (1 - dx) + bottom_right_value * dy * dx +
                   top_right_value * (1 - dy) * dx + bottom_left_value * dy * (1 - dx);
    return static_cast<DType>(result);
  }

 private:
  Plan<SrcExp, DType> src_;
  const float start_y_;
  const float start_x_;
  const float step_y_;
  const float step_x_;
  const index_t src_height_;
  const index_t src_width_;
  const index_t out_height_;
  const int pad_mode_;
  const DType pad_value_;
};

用户接口

用户在使用的时候，调用 resize 之后实际上只是构建了一个 ResizeExp 表达式，而没有发生任何有用的计算，只有在 resize “赋值”给一个 Tensor 的时候计算才真正发生。所谓赋值不仅仅包括=，也包括+=，-=等，具体可以参考RValueExp中的定义 Tensor 的定义

1
2
3
4
5
6
7
8
template <typename SrcExp, typename DType, int etype>
inline ResizeExp<SrcExp, DType, ExpInfo<SrcExp>::kDim> resize(const Exp<SrcExp, DType, etype>& src,
                                                              index_t out_height, index_t out_width,
                                                              int pad_mode = resize_pad::kEdge,
                                                              DType pad_value = 0) {
  return ResizeExp<SrcExp, DType, ExpInfo<SrcExp>::kDim>(src.self(), out_height, out_width,
                                                         pad_mode, pad_value);
}