fix the precise roi poop op test=develop (PaddlePaddle#20126)

* fix the precise roi poop op test=develop
add roi backward implementation, fix the output-channel

Author: wopeizl

paddle/fluid/API.spec

@@ -290,7 +290,7 @@ paddle.fluid.layers.shuffle_channel (ArgSpec(args=['x', 'group', 'name'], vararg
 paddle.fluid.layers.temporal_shift (ArgSpec(args=['x', 'seg_num', 'shift_ratio', 'name'], varargs=None, keywords=None, defaults=(0.25, None)), ('document', 'd5945431cdcae3cda21914db5bbf383e'))
 paddle.fluid.layers.py_func (ArgSpec(args=['func', 'x', 'out', 'backward_func', 'skip_vars_in_backward_input'], varargs=None, keywords=None, defaults=(None, None)), ('document', '231f91231430f5dae2b757df22317c67'))
 paddle.fluid.layers.psroi_pool (ArgSpec(args=['input', 'rois', 'output_channels', 'spatial_scale', 'pooled_height', 'pooled_width', 'name'], varargs=None, keywords=None, defaults=(None,)), ('document', '9bf0cc6b0717010b8ceec5dc2541d566'))
-paddle.fluid.layers.prroi_pool (ArgSpec(args=['input', 'rois', 'output_channels', 'spatial_scale', 'pooled_height', 'pooled_width', 'name'], varargs=None, keywords=None, defaults=(1.0, 1, 1, None)), ('document', '454c7ea8c73313dd41513929d7526303'))
+paddle.fluid.layers.prroi_pool (ArgSpec(args=['input', 'rois', 'spatial_scale', 'pooled_height', 'pooled_width', 'name'], varargs=None, keywords=None, defaults=(1.0, 1, 1, None)), ('document', '466be691ac4c1cd7b88fccb40846afce'))
 paddle.fluid.layers.teacher_student_sigmoid_loss (ArgSpec(args=['input', 'label', 'soft_max_up_bound', 'soft_max_lower_bound'], varargs=None, keywords=None, defaults=(15.0, -15.0)), ('document', 'b0e07aa41caae04b07a8e8217cc96020'))
 paddle.fluid.layers.huber_loss (ArgSpec(args=['input', 'label', 'delta'], varargs=None, keywords=None, defaults=None), ('document', '9d93ee81f7a3e526d68bb280bc695d6c'))
 paddle.fluid.layers.kldiv_loss (ArgSpec(args=['x', 'target', 'reduction', 'name'], varargs=None, keywords=None, defaults=('mean', None)), ('document', '45f3ebbcb766fca84cb2fe6307086573'))

paddle/fluid/operators/prroi_pool_op.cc

@@ -43,12 +43,6 @@ class PRROIPoolOpMaker : public framework::OpProtoAndCheckerMaker {
               "(Tensor), "
               "the output of PRROIPoolOp is a 4-D Tensor with shape "
               "(num_rois, output_channels, pooled_h, pooled_w).");
-    AddAttr<int>(
-        "output_channels",
-        "(int), "
-        "the number of channels of the output feature map. "
-        "For a task of C classes of objects, output_channels should be "
-        "(C + 1) for classification only.");
     AddAttr<float>("spatial_scale",
                    "(float, default 1.0), "
                    "Multiplicative spatial scale factor "
@@ -100,28 +94,18 @@ class PRROIPoolOp : public framework::OperatorWithKernel {
 
     int pooled_height = ctx->Attrs().Get<int>("pooled_height");
     int pooled_width = ctx->Attrs().Get<int>("pooled_width");
-    int output_channels = ctx->Attrs().Get<int>("output_channels");
     float spatial_scale = ctx->Attrs().Get<float>("spatial_scale");
 
-    PADDLE_ENFORCE_EQ(
-        input_dims[1], output_channels * pooled_height * pooled_width,
-        "the channel of X(%d) should be equal to the product of "
-        "output_channels(%d), pooled_height(%d) and pooled_width(%d)",
-        input_dims[1], output_channels, pooled_height, pooled_width);
-
     PADDLE_ENFORCE_GT(pooled_height, 0,
                       "The pooled output height must be greater than 0");
     PADDLE_ENFORCE_GT(pooled_width, 0,
                       "The pooled output width must be greater than 0");
-    PADDLE_ENFORCE_GT(output_channels, 1,
-                      "The pooled output channels must greater than 1");
     PADDLE_ENFORCE_GT(spatial_scale, 0.0f,
                       "The spatial scale must greater than 0.");
 
     auto out_dims = input_dims;
     out_dims[0] = rois_dims[0];
-    out_dims[1] =
-        output_channels;  // input_dims[1] / (pooled_height * pooled_width);
+    out_dims[1] = input_dims[1];
     out_dims[2] = pooled_height;
     out_dims[3] = pooled_width;
     ctx->SetOutputDim("Out", out_dims);
@@ -145,6 +129,7 @@ class PRROIPoolGradOp : public framework::OperatorWithKernel {
     PADDLE_ENFORCE_EQ(ctx->HasOutput(framework::GradVarName("X")), true,
                       "The gradient of X should not be null.");
     ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
+    ctx->SetOutputDim(framework::GradVarName("ROIs"), ctx->GetInputDim("ROIs"));
   }
 
  protected:
@@ -164,9 +149,11 @@ class PRROIPoolGradDescMaker : public framework::SingleGradOpDescMaker {
     std::unique_ptr<framework::OpDesc> op(new framework::OpDesc());
     op->SetType("prroi_pool_grad");
     op->SetInput("X", Input("X"));
+    op->SetInput("Out", Output("Out"));
     op->SetInput("ROIs", Input("ROIs"));
     op->SetInput(framework::GradVarName("Out"), OutputGrad("Out"));
     op->SetOutput(framework::GradVarName("X"), InputGrad("X"));
+    op->SetOutput(framework::GradVarName("ROIs"), InputGrad("ROIs"));
     op->SetAttrMap(Attrs());
     return op;
   }

paddle/fluid/operators/prroi_pool_op.cu

@@ -40,6 +40,11 @@ DEVICE void PrRoIPoolingDistributeDiffCUDA(T* diff, const T top_diff,
   }
 }
 
+template <typename T>
+DEVICE void GPUAccumulateRois(T* offset, T data) {
+  paddle::platform::CudaAtomicAdd(offset, data);
+}
+
 template <typename T>
 __global__ void GPUPRROIPoolForward(
     const int nthreads, const T* input_data, const T* input_rois,
@@ -78,7 +83,7 @@ __global__ void GPUPRROIPoolForward(
     T win_end_h = win_start_h + bin_size_h;
 
     T win_size = max(static_cast<T>(0.0), bin_size_w * bin_size_h);
-    int input_channel = (c * pooled_height + ph) * pooled_width + pw;
+    int input_channel = c;
     const T* offset_input_data =
         input_data +
         (roi_batch_id * input_channels + input_channel) * height * width;
@@ -110,10 +115,12 @@ __global__ void GPUPRROIPoolForward(
 
 template <typename T>
 __global__ void GPUPRROIPoolBackward(
-    const int nthreads, const T* input_rois, const T* output_grad_data,
-    const float spatial_scale, const int input_channels, const int height,
-    const int width, const int output_channels, const int pooled_height,
-    const int pooled_width, const int* rois_batch_id_data, T* input_grad_data) {
+    const int nthreads, const T* in_data, const T* input_rois,
+    const T* output_grad_data, const float spatial_scale,
+    const int input_channels, const int height, const int width,
+    const int output_channels, const int pooled_height, const int pooled_width,
+    const int* rois_batch_id_data, T* input_grad_data, const T* out_data,
+    T* input_roi_grad_data) {
   int index = blockIdx.x * blockDim.x + threadIdx.x;
   int offset = blockDim.x * gridDim.x;
   for (int i = index; i < nthreads; i += offset) {
@@ -125,7 +132,7 @@ __global__ void GPUPRROIPoolBackward(
 
     // set roi_batch_id
     int roi_batch_id = rois_batch_id_data[n];
-    int input_channel = (c * pooled_height + ph) * pooled_width + pw;
+    int input_channel = c;
     int input_offset =
         (roi_batch_id * input_channels + input_channel) * height * width;
     T* offset_input_grad_data = input_grad_data + input_offset;
@@ -137,6 +144,7 @@ __global__ void GPUPRROIPoolBackward(
     T roi_start_h = static_cast<T>(offset_input_rois[1]) * spatial_scale;
     T roi_end_w = static_cast<T>(offset_input_rois[2]) * spatial_scale;
     T roi_end_h = static_cast<T>(offset_input_rois[3]) * spatial_scale;
+    T* offset_input_roi_grad_data = input_roi_grad_data + n * 4;
 
     T roi_width = max(roi_end_w - roi_start_w, static_cast<T>(0.0));
     T roi_height = max(roi_end_h - roi_start_h, static_cast<T>(0.0));
@@ -171,6 +179,16 @@ __global__ void GPUPRROIPoolBackward(
             height, width, PrRoIPoolingDistributeDiffCUDA<T>);
       }
     }
+
+    const T* offset_out_data = out_data + i;
+    const T* offset_in_data = in_data + input_offset;
+    PrRoIPoolingCoorBackward(
+        s_w, e_w, s_h, e_h, width, height, win_start_w, win_start_h, win_end_w,
+        win_end_h, pw, ph, pooled_width, pooled_height, win_size, spatial_scale,
+        offset_in_data, offset_out_data, offset_input_grad_data,
+        offset_input_roi_grad_data, GPUAccumulateRois<T>,
+        [](const T x, const T y) { return max(x, y); },
+        [](const T x, const T y) { return min(x, y); });
   }
 }
 
@@ -184,20 +202,15 @@ class GPUPRROIPoolOpKernel : public framework::OpKernel<T> {
 
     auto pooled_height = ctx.Attr<int>("pooled_height");
     auto pooled_width = ctx.Attr<int>("pooled_width");
-    auto output_channels = ctx.Attr<int>("output_channels");
     auto spatial_scale = ctx.Attr<float>("spatial_scale");
 
     auto in_dims = in->dims();
     int batch_size = in_dims[0];
     int input_channels = in_dims[1];
+    auto output_channels = input_channels;
     int height = in_dims[2];
     int width = in_dims[3];
 
-    PADDLE_ENFORCE_EQ(input_channels,
-                      output_channels * pooled_height * pooled_width,
-                      "the channels of input X should equal the product of "
-                      "output_channels x pooled_height x pooled_width");
-
     int rois_num = rois->dims()[0];
     if (rois_num == 0) return;
 
@@ -245,17 +258,20 @@ class GPUPRROIPoolGradOpKernel : public framework::OpKernel<T> {
   void Compute(const framework::ExecutionContext& ctx) const override {
     auto* in = ctx.Input<Tensor>("X");
     auto* rois = ctx.Input<LoDTensor>("ROIs");
+    auto* out = ctx.Input<framework::Tensor>("Out");
 
     auto* output_grad = ctx.Input<Tensor>(framework::GradVarName("Out"));
     auto* input_grad = ctx.Output<Tensor>(framework::GradVarName("X"));
+    auto* input_roi_grad =
+        ctx.Output<LoDTensor>(framework::GradVarName("ROIs"));
 
     auto pooled_height = ctx.Attr<int>("pooled_height");
     auto pooled_width = ctx.Attr<int>("pooled_width");
-    auto output_channels = ctx.Attr<int>("output_channels");
     auto spatial_scale = ctx.Attr<float>("spatial_scale");
 
     int rois_num = rois->dims()[0];
     int input_channels = in->dims()[1];
+    auto output_channels = input_channels;
     int height = in->dims()[2];
     int width = in->dims()[3];
 
@@ -280,6 +296,8 @@ class GPUPRROIPoolGradOpKernel : public framework::OpKernel<T> {
       input_grad->mutable_data<T>(ctx.GetPlace());
       math::SetConstant<DeviceContext, T> set_zero;
       set_zero(ctx.cuda_device_context(), input_grad, static_cast<T>(0));
+      input_roi_grad->mutable_data<T>(ctx.GetPlace());
+      set_zero(ctx.cuda_device_context(), input_roi_grad, static_cast<T>(0));
 
       int output_grad_size = output_grad->numel();
       int blocks = NumBlocks(output_grad_size);
@@ -288,10 +306,12 @@ class GPUPRROIPoolGradOpKernel : public framework::OpKernel<T> {
       if (output_grad_size > 0) {
         GPUPRROIPoolBackward<
             T><<<blocks, threads, 0, ctx.cuda_device_context().stream()>>>(
-            output_grad_size, rois->data<T>(), output_grad->data<T>(),
-            spatial_scale, input_channels, height, width, output_channels,
-            pooled_height, pooled_width, rois_batch_id_list_gpu.data<int>(),
-            input_grad->mutable_data<T>(ctx.GetPlace()));
+            output_grad_size, in->data<T>(), rois->data<T>(),
+            output_grad->data<T>(), spatial_scale, input_channels, height,
+            width, output_channels, pooled_height, pooled_width,
+            rois_batch_id_list_gpu.data<int>(),
+            input_grad->mutable_data<T>(ctx.GetPlace()), out->data<T>(),
+            input_roi_grad->mutable_data<T>(ctx.GetPlace()));
       }
     }
   }