channel_unpool.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#ifndef MSHADOW_EXTENSION_CHANNEL_UNPOOL_H_
#define MSHADOW_EXTENSION_CHANNEL_UNPOOL_H_
#include <algorithm>
#include "../extension.h"
namespace mshadow {
namespace expr {
template<typename Reducer, typename SrcExp, typename DType, int srcdim>
struct ChannelUnpoolingExp:
      public MakeTensorExp<ChannelUnpoolingExp<Reducer, SrcExp, DType, srcdim>,
                           SrcExp, srcdim, DType> {
  const SrcExp &data_src_;
  const SrcExp &data_pooled_;
  const SrcExp &grad_pooled_;
  index_t pchannel_;
  index_t nsize_;
  index_t kstride_;
  index_t pad_;
  ChannelUnpoolingExp(const SrcExp &data_src,
               const SrcExp &data_pooled,
               const SrcExp &grad_pooled,
               index_t nsize, index_t kstride, index_t pad)
      : data_src_(data_src), data_pooled_(data_pooled),
        grad_pooled_(grad_pooled),
        nsize_(nsize), kstride_(kstride), pad_(pad) {
    Shape<srcdim> pshape = ShapeCheck<srcdim, SrcExp>::Check(grad_pooled);
    typedef ShapeCheck<srcdim, SrcExp> ShapeCheckSrcDimSrcExp;
    CHECK_EQ(pshape, ShapeCheckSrcDimSrcExp::Check(data_pooled))
      << "ChannelUnPoolingExp: data and grad shape mismatch";
    Shape<srcdim> sshape = ShapeCheck<srcdim, SrcExp>::Check(data_src);
    for (int k = 0; k < srcdim; ++k) {
      if (k == 1) {
        continue;
      }
      CHECK_EQ(pshape[k], sshape[k])
        << "ChannelUnPoolingExp: pooled tensor and src tensor shape mismatch"
        << pshape[k]
        << " vs "
        << sshape[k];
    }
    pchannel_ = pshape[1];
    this->shape_ = sshape;
  }
};
template<typename Reducer, typename SrcExp, typename DType, int etype>
inline ChannelUnpoolingExp<Reducer, SrcExp, DType, ExpInfo<SrcExp>::kDim>
ch_unpool(const Exp<SrcExp, DType, etype> &data_src,
       const Exp<SrcExp, DType, etype> &data_pooled,
       const Exp<SrcExp, DType, etype> &grad_pooled,
      index_t nsize, index_t stride, index_t pad) {
  TypeCheckPass<ExpInfo<SrcExp>::kDim >= 3>
      ::Error_Expression_Does_Not_Meet_Dimension_Req();
  return ChannelUnpoolingExp<Reducer, SrcExp, DType, ExpInfo<SrcExp>::kDim>
        (data_src.self(), data_pooled.self(), grad_pooled.self(), nsize, stride, pad);
}

template<typename Reducer, typename SrcExp, typename DType, int etype>
inline ChannelUnpoolingExp<Reducer, SrcExp, DType, ExpInfo<SrcExp>::kDim>
ch_unpool(const Exp<SrcExp, DType, etype> &data_src,
       const Exp<SrcExp, DType, etype> &data_pooled,
       const Exp<SrcExp, DType, etype> &grad_pooled, index_t nsize) {
  return ch_unpool(data_src, data_pooled, grad_pooled, nsize, 1, nsize / 2);
}


//----------------------
// Execution plan
//----------------------
template<typename Reducer, typename SrcExp, typename DType, int srcdim>
struct Plan<ChannelUnpoolingExp<Reducer, SrcExp, DType, srcdim>, DType> {
 public:
  explicit Plan(const ChannelUnpoolingExp<Reducer, SrcExp, DType, srcdim> &e)
      : data_src_(e.data_src_), data_pooled_(e.data_pooled_),
        grad_pooled_(e.grad_pooled_), channel_(e.shape_[srcdim - 3]),
        height_(e.shape_[srcdim - 2]), pchannel_(e.pchannel_),
        hnsize_(e.nsize_), stride_(e.kstride_), pad_(e.pad_) {}
  MSHADOW_XINLINE DType Eval(index_t i, index_t j) const {
    using namespace std;
    const DType vsrc = data_src_.Eval(i, j);
    const index_t y = i % height_;
    i /= height_;
    const index_t c = i % channel_;
    const index_t n = i / channel_;
    const index_t x = j;
    const index_t cstart = c < hnsize_ - pad_ ? 0
                        : (c - (hnsize_ - pad_) + stride_) / stride_;
    const index_t cend = min((c + pad_ + stride_) / stride_, channel_);
    DType val = static_cast<DType>(0);
    for (index_t cc = cstart; cc < cend; ++cc) {
      val += Reducer::PartialGrad(vsrc,
                                  data_pooled_.Eval((n * pchannel_ + cc) * height_ + y, x)) *
                                  grad_pooled_.Eval((n * pchannel_ + cc) * height_ + y, x);
    }
    return val;
  }

 private:
  Plan<SrcExp, DType> data_src_, data_pooled_, grad_pooled_;
  const index_t channel_, height_, pchannel_, hnsize_, stride_, pad_;
};
}  // namespace expr
}  // namespace mshadow
#endif  // MSHADOW_EXTENSION_CHANNEL_UNPOOL_H_