roi_pool_op.h 9.4 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
W
wanghaox 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15

Licensed 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. */

#pragma once
16 17
#include <algorithm>
#include <limits>
F
FDInSky 已提交
18
#include <vector>
Y
Yi Wang 已提交
19
#include "paddle/fluid/framework/op_registry.h"
F
FDInSky 已提交
20
#include "paddle/fluid/memory/memcpy.h"
Y
Yi Wang 已提交
21
#include "paddle/fluid/operators/math/math_function.h"
W
wanghaox 已提交
22 23 24 25

namespace paddle {
namespace operators {

26 27
static constexpr int kROISize = 4;

Q
QI JUN 已提交
28
template <typename DeviceContext, typename T>
W
wanghaox 已提交
29
class CPUROIPoolOpKernel : public framework::OpKernel<T> {
W
wanghaox 已提交
30 31
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
W
wanghaox 已提交
32
    auto* in = ctx.Input<framework::Tensor>("X");
33
    auto* rois = ctx.Input<framework::LoDTensor>("ROIs");
W
wanghaox 已提交
34 35
    auto* out = ctx.Output<framework::Tensor>("Out");
    auto* argmax = ctx.Output<framework::Tensor>("Argmax");
W
wanghaox 已提交
36 37 38 39 40 41 42 43 44 45 46 47 48 49 50

    auto pooled_height = ctx.Attr<int>("pooled_height");
    auto pooled_width = ctx.Attr<int>("pooled_width");
    auto spatial_scale = ctx.Attr<float>("spatial_scale");

    auto in_dims = in->dims();
    int batch_size = in_dims[0];
    int channels = in_dims[1];
    int height = in_dims[2];
    int width = in_dims[3];
    int rois_num = rois->dims()[0];

    auto in_stride = framework::stride(in_dims);
    auto argmax_stride = framework::stride(argmax->dims());
    auto roi_stride = framework::stride(rois->dims());
W
wanghaox 已提交
51
    auto out_stride = framework::stride(out->dims());
W
wanghaox 已提交
52 53 54

    const T* input_data = in->data<T>();

55 56 57 58 59
    framework::Tensor roi_batch_id_list;
    roi_batch_id_list.Resize({rois_num});
    int* roi_batch_id_data =
        roi_batch_id_list.mutable_data<int>(ctx.GetPlace());

F
FDInSky 已提交
60
    int rois_batch_size;
61 62 63
    if (ctx.HasInput("RoisNum")) {
      auto* rois_num_t = ctx.Input<framework::Tensor>("RoisNum");
      rois_batch_size = rois_num_t->numel();
F
FDInSky 已提交
64
      PADDLE_ENFORCE_EQ(
65
          rois_batch_size, batch_size,
66 67
          platform::errors::InvalidArgument("The rois_batch_size and imgs "
                                            "batch_size must be the same."));
68 69 70 71
      auto* rois_num_data = rois_num_t->data<int>();
      int start = 0;
      for (int n = 0; n < rois_batch_size; ++n) {
        for (int i = start; i < start + rois_num_data[n]; ++i) {
F
FDInSky 已提交
72 73
          roi_batch_id_data[i] = n;
        }
74
        start += rois_num_data[n];
F
FDInSky 已提交
75 76 77 78 79 80
      }
    } else {
      auto rois_lod = rois->lod().back();
      rois_batch_size = rois_lod.size() - 1;
      PADDLE_ENFORCE_EQ(
          rois_batch_size, batch_size,
81 82
          platform::errors::InvalidArgument("The rois_batch_size and imgs "
                                            "batch_size must be the same."));
F
FDInSky 已提交
83
      int rois_num_with_lod = rois_lod[rois_batch_size];
84 85 86 87
      PADDLE_ENFORCE_EQ(
          rois_num, rois_num_with_lod,
          platform::errors::InvalidArgument("The rois_num from input "
                                            "and lod must be the same."));
F
FDInSky 已提交
88 89 90 91
      for (int n = 0; n < rois_batch_size; ++n) {
        for (size_t i = rois_lod[n]; i < rois_lod[n + 1]; ++i) {
          roi_batch_id_data[i] = n;
        }
92
      }
W
wanghaox 已提交
93 94
    }

95 96 97
    T* output_data = out->mutable_data<T>(ctx.GetPlace());
    int64_t* argmax_data = argmax->mutable_data<int64_t>(ctx.GetPlace());

98
    const T* rois_data = rois->data<T>();
W
wanghaox 已提交
99
    for (int n = 0; n < rois_num; ++n) {
100 101 102 103 104
      int roi_batch_id = roi_batch_id_data[n];
      int roi_start_w = round(rois_data[0] * spatial_scale);
      int roi_start_h = round(rois_data[1] * spatial_scale);
      int roi_end_w = round(rois_data[2] * spatial_scale);
      int roi_end_h = round(rois_data[3] * spatial_scale);
W
wanghaox 已提交
105 106 107 108 109 110 111 112 113 114

      // Force malformed ROIs to be 1x1
      int roi_height = std::max(roi_end_h - roi_start_h + 1, 1);
      int roi_width = std::max(roi_end_w - roi_start_w + 1, 1);

      const float bin_size_h =
          static_cast<float>(roi_height) / static_cast<float>(pooled_height);
      const float bin_size_w =
          static_cast<float>(roi_width) / static_cast<float>(pooled_width);

W
wanghaox 已提交
115
      const T* batch_data = input_data + roi_batch_id * in_stride[0];
W
wanghaox 已提交
116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140

      for (int c = 0; c < channels; ++c) {
        for (int ph = 0; ph < pooled_height; ++ph) {
          for (int pw = 0; pw < pooled_width; ++pw) {
            //  Compute pooling region for this output unit:
            //  start (included) = floor(ph * roi_height / pooled_height_)
            //  end (excluded) = ceil((ph + 1) * roi_height / pooled_height_)
            int hstart =
                static_cast<int>(floor(static_cast<float>(ph) * bin_size_h));
            int wstart =
                static_cast<int>(floor(static_cast<float>(pw) * bin_size_w));
            int hend =
                static_cast<int>(ceil(static_cast<float>(ph + 1) * bin_size_h));
            int wend =
                static_cast<int>(ceil(static_cast<float>(pw + 1) * bin_size_w));

            hstart = std::min(std::max(hstart + roi_start_h, 0), height);
            hend = std::min(std::max(hend + roi_start_h, 0), height);
            wstart = std::min(std::max(wstart + roi_start_w, 0), width);
            wend = std::min(std::max(wend + roi_start_w, 0), width);

            const int pool_index = ph * pooled_width + pw;

            // Define an empty pooling region to be zero
            bool is_empty = (hend <= hstart) || (wend <= wstart);
W
wanghaox 已提交
141
            output_data[pool_index] =
W
wanghaox 已提交
142 143
                is_empty ? 0 : -std::numeric_limits<T>::max();
            argmax_data[pool_index] = -1;
W
wanghaox 已提交
144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167

            for (int h = hstart; h < hend; ++h) {
              for (int w = wstart; w < wend; ++w) {
                const int index = h * width + w;
                if (batch_data[index] > output_data[pool_index]) {
                  output_data[pool_index] = batch_data[index];
                  argmax_data[pool_index] = index;
                }
              }
            }
          }
        }

        batch_data += in_stride[1];
        output_data += out_stride[1];
        argmax_data += argmax_stride[1];
      }
      // Increment ROI data pointer
      rois_data += roi_stride[0];
    }
    return;
  }
};

Q
QI JUN 已提交
168
template <typename DeviceContext, typename T>
W
wanghaox 已提交
169
class CPUROIPoolGradOpKernel : public framework::OpKernel<T> {
W
wanghaox 已提交
170 171
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
W
wanghaox 已提交
172
    auto* in = ctx.Input<framework::Tensor>("X");
173
    auto* rois = ctx.Input<framework::LoDTensor>("ROIs");
W
wanghaox 已提交
174
    auto* argmax = ctx.Input<framework::Tensor>("Argmax");
W
wanghaox 已提交
175
    auto* out_grad =
W
wanghaox 已提交
176
        ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
G
guosheng 已提交
177
    auto* in_grad = ctx.Output<framework::Tensor>(framework::GradVarName("X"));
W
wanghaox 已提交
178 179 180 181

    auto pooled_height = ctx.Attr<int>("pooled_height");
    auto pooled_width = ctx.Attr<int>("pooled_width");

G
guosheng 已提交
182
    if (in_grad) {
183 184 185 186 187 188
      int rois_num = rois->dims()[0];
      framework::Tensor roi_batch_id_list;
      roi_batch_id_list.Resize({rois_num});
      int* roi_batch_id_data =
          roi_batch_id_list.mutable_data<int>(ctx.GetPlace());

F
FDInSky 已提交
189
      int rois_batch_size;
190 191 192 193 194 195 196
      if (ctx.HasInput("RoisNum")) {
        auto* rois_num_t = ctx.Input<framework::Tensor>("RoisNum");
        rois_batch_size = rois_num_t->numel();
        auto* rois_num_data = rois_num_t->data<int>();
        int start = 0;
        for (int n = 0; n < rois_batch_size; ++n) {
          for (int i = start; i < start + rois_num_data[n]; ++i) {
F
FDInSky 已提交
197 198
            roi_batch_id_data[i] = n;
          }
199
          start += rois_num_data[n];
F
FDInSky 已提交
200 201 202 203 204 205 206 207
        }
      } else {
        auto rois_lod = rois->lod().back();
        rois_batch_size = rois_lod.size() - 1;
        for (int n = 0; n < rois_batch_size; ++n) {
          for (size_t i = rois_lod[n]; i < rois_lod[n + 1]; ++i) {
            roi_batch_id_data[i] = n;
          }
208 209 210
        }
      }

211
      const T* rois_data = rois->data<T>();
G
guosheng 已提交
212 213 214
      const T* out_grad_data = out_grad->data<T>();
      const int64_t* argmax_data = argmax->data<int64_t>();
      T* in_grad_data = in_grad->mutable_data<T>(ctx.GetPlace());
Q
QI JUN 已提交
215 216 217
      math::SetConstant<DeviceContext, T> set_zero;
      set_zero(ctx.template device_context<DeviceContext>(), in_grad,
               static_cast<T>(0));
W
wanghaox 已提交
218

G
guosheng 已提交
219 220 221 222
      auto in_stride = framework::stride(in->dims());
      auto argmax_stride = framework::stride(argmax->dims());
      auto roi_stride = framework::stride(rois->dims());
      auto out_stride = framework::stride(out_grad->dims());
W
wanghaox 已提交
223

G
guosheng 已提交
224
      int channels = in->dims()[1];
W
wanghaox 已提交
225

G
guosheng 已提交
226
      for (int n = 0; n < rois_num; ++n) {
227
        int roi_batch_idx = roi_batch_id_data[n];
G
guosheng 已提交
228
        T* batch_grad_data = in_grad_data + roi_batch_idx * in_stride[0];
W
wanghaox 已提交
229 230 231
        for (int c = 0; c < channels; ++c) {
          for (int ph = 0; ph < pooled_height; ++ph) {
            for (int pw = 0; pw < pooled_width; ++pw) {
G
guosheng 已提交
232
              int pool_index = ph * pooled_width + pw;
W
wanghaox 已提交
233
              if (argmax_data[pool_index] >= 0) {
G
guosheng 已提交
234
                auto index = argmax_data[pool_index];
W
wanghaox 已提交
235 236 237 238
                batch_grad_data[index] += out_grad_data[pool_index];
              }
            }
          }
G
guosheng 已提交
239 240 241
          batch_grad_data += in_stride[1];
          out_grad_data += out_stride[1];
          argmax_data += argmax_stride[1];
W
wanghaox 已提交
242
        }
G
guosheng 已提交
243
        rois_data += roi_stride[0];
W
wanghaox 已提交
244 245 246 247 248 249 250
      }
    }
  }
};

}  // namespace operators
}  // namespace paddle