unpack_patch2col.h

Go to the documentation of this file.

#ifndef MSHADOW_EXTENSION_UNPACK_PATCH2COL_H_
#define MSHADOW_EXTENSION_UNPACK_PATCH2COL_H_
#include "../extension.h"
namespace mshadow {
namespace expr {
template<typename SrcExp, typename DType, int srcdim>
struct UnpackPatchToColXExp:
      public MakeTensorExp<UnpackPatchToColXExp<SrcExp, DType, srcdim>,
                           SrcExp, 2, DType>{
  const SrcExp &img_;
  index_t psize_y_;
  index_t psize_x_;
  index_t pstride_y_;
  index_t pstride_x_;
  index_t pdilate_y_;
  index_t pdilate_x_;
  index_t i_channel_;
  index_t i_height_;
  index_t i_width_;
  UnpackPatchToColXExp(const SrcExp &img,
                       index_t psize_y,
                       index_t psize_x,
                       index_t pstride_y,
                       index_t pstride_x,
                       index_t pdilate_y,
                       index_t pdilate_x)
      : img_(img), psize_y_(psize_y), psize_x_(psize_x),
      pstride_y_(pstride_y), pstride_x_(pstride_x),
      pdilate_y_(pdilate_y), pdilate_x_(pdilate_x){
    Shape<srcdim> imshape = ShapeCheck<srcdim, SrcExp>::Check(img_);
    CHECK(imshape[srcdim - 1] >= psize_x && imshape[srcdim - 2] >= psize_y)
      << "UnpackPatchToCol:image shape smaller than patch size";
    this->i_channel_ = imshape[srcdim - 3];
    this->i_height_  = imshape[srcdim - 2];
    this->i_width_   = imshape[srcdim - 1];
    // calculate number of batches
    const index_t num = imshape.ProdShape(0, srcdim - 3);
    const index_t o_height = (i_height_ -
        (pdilate_y * (psize_y - 1) + 1)) / pstride_y + 1;
    const index_t o_width  = (i_width_  -
        (pdilate_x * (psize_x - 1) + 1)) / pstride_x + 1;
    this->shape_[1] = o_height * o_width * num;
    this->shape_[0] = psize_y * psize_x * i_channel_;
  }
};

template<typename SrcExp, typename DType, int etype>
inline UnpackPatchToColXExp<SrcExp, DType, ExpInfo<SrcExp>::kDim>
unpack_patch2col(const Exp<SrcExp, DType, etype> &img,
                 index_t psize_y, index_t psize_x, index_t pstride, index_t pdilate) {
  TypeCheckPass<ExpInfo<SrcExp>::kDim >= 3>
      ::Error_Expression_Does_Not_Meet_Dimension_Req();
  return UnpackPatchToColXExp<SrcExp, DType, ExpInfo<SrcExp>::kDim>
      (img.self(), psize_y, psize_x, pstride, pstride, pdilate, pdilate);
}

template<typename SrcExp, typename DType, int etype>
inline UnpackPatchToColXExp<SrcExp, DType, ExpInfo<SrcExp>::kDim>
unpack_patch2col(const Exp<SrcExp, DType, etype> &img,
                 index_t psize_y, index_t psize_x, index_t pstride_y_, index_t pstride_x_,
                 index_t pdilate_y_, index_t pdilate_x_) {
  TypeCheckPass<ExpInfo<SrcExp>::kDim >= 3>
      ::Error_Expression_Does_Not_Meet_Dimension_Req();
  return UnpackPatchToColXExp<SrcExp, DType, ExpInfo<SrcExp>::kDim>
      (img.self(), psize_y, psize_x, pstride_y_, pstride_x_, pdilate_y_, pdilate_x_);
}
//----------------------
// Execution plan
//----------------------
template<typename SrcExp, typename DType, int srcdim>
struct Plan<UnpackPatchToColXExp<SrcExp, DType, srcdim>, DType> {
 public:
  explicit Plan(const UnpackPatchToColXExp<SrcExp, DType, srcdim> &e)
      :src_(MakePlan(e.img_)),
       psize_y_(e.psize_y_), psize_x_(e.psize_x_),
       pstride_y_(e.pstride_y_), pstride_x_(e.pstride_x_),
       i_channel_(e.i_channel_), pdilate_y_(e.pdilate_y_), pdilate_x_(e.pdilate_x_),
       i_height_(e.i_height_), i_width_(e.i_width_),
       o_height_((i_height_ - (pdilate_y_ * (psize_y_ - 1) + 1)) / pstride_y_ + 1),
       o_width_((i_width_ - (pdilate_x_ * (psize_x_ - 1) + 1)) / pstride_x_ + 1) {}
  MSHADOW_XINLINE DType Eval(index_t i, index_t j) const {
    const index_t x_offset = i % psize_x_ * pdilate_x_;
    const index_t idivp    = i / psize_x_;
    const index_t y_offset = idivp % psize_y_ * pdilate_y_;
    const index_t c = idivp / psize_y_;
    const index_t x = (j % o_width_) * pstride_x_ + x_offset;
    const index_t jdivw = j / o_width_;
    const index_t y = (jdivw % o_height_) * pstride_y_ + y_offset;
    const index_t n = jdivw / o_height_;

    if (x < i_width_ && y < i_height_) {
      return src_.Eval((n * i_channel_  + c) * i_height_ + y, x);
    } else {
      return DType(0.0f);
    }
  }

 private:
  Plan<SrcExp, DType> src_;
  const index_t psize_y_, psize_x_, pstride_y_, pstride_x_, i_channel_;
  const index_t pdilate_y_, pdilate_x_;
  const index_t i_height_, i_width_, o_height_, o_width_;
};
}  // namespace expr
}  // namespace mshadow
#endif  // MSHADOW_EXTENSION_UNPACK_PATCH2COL_H_