Add grad grad for AvgPool2D (#35388)

* add pool2d grad grad

* dbg

* add unittest

* update format

* add more unittests

* dbg

paddle/fluid/operators/pool_cudnn_op.cu.cc

@@ -505,6 +505,20 @@ class PoolCUDNNGradOpKernel : public framework::OpKernel<T> {
   }
 };
 
+template <typename T>
+class PoolCUDNNGradGradOpKernel : public PoolCUDNNOpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext &ctx) const override {
+    std::string pooling_type = ctx.Attr<std::string>("pooling_type");
+    if (pooling_type == "max") {
+      PADDLE_THROW(platform::errors::InvalidArgument(
+          "Pool op grad grad only supports avgpool."));
+    } else {
+      PoolCUDNNOpKernel<T>::Compute(ctx);
+    }
+  }
+};
+
 }  // namespace operators
 }  // namespace paddle
 
@@ -534,6 +548,10 @@ REGISTER_OP_KERNEL(pool2d_grad, CUDNN, plat::CUDAPlace,
                    ops::PoolCUDNNGradOpKernel<float>,
                    ops::PoolCUDNNGradOpKernel<double>,
                    ops::PoolCUDNNGradOpKernel<plat::float16>);
+REGISTER_OP_KERNEL(pool2d_grad_grad, CUDNN, plat::CUDAPlace,
+                   ops::PoolCUDNNGradGradOpKernel<float>,
+                   ops::PoolCUDNNGradGradOpKernel<double>,
+                   ops::PoolCUDNNGradGradOpKernel<plat::float16>);
 
 REGISTER_OP_KERNEL(pool3d, CUDNN, plat::CUDAPlace,
                    ops::PoolCUDNNOpKernel<float>,

paddle/fluid/operators/pool_op.cc

@@ -469,6 +469,20 @@ Example:
 )DOC");
 }
 
+template <typename T>
+class Pool2dOpGradGradMaker : public framework::SingleGradOpMaker<T> {
+ public:
+  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
+
+ protected:
+  void Apply(GradOpPtr<T> grad_op) const override {
+    grad_op->SetType("pool2d_grad_grad");
+    grad_op->SetInput("X", this->OutputGrad(framework::GradVarName("X")));
+    grad_op->SetOutput("Out", this->InputGrad(framework::GradVarName("Out")));
+    grad_op->SetAttrMap(this->Attrs());
+  }
+};
+
 class PoolOpInferVarType : public framework::PassInDtypeAndVarTypeToOutput {
  protected:
   std::unordered_map<std::string, std::string>& GetInputOutputWithSameType()
@@ -687,7 +701,10 @@ REGISTER_OPERATOR(
     pool2d, ops::PoolOp, ops::Pool2dOpMaker, ops::PoolOpInferVarType,
     paddle::framework::DefaultGradOpMaker<paddle::framework::OpDesc, true>,
     paddle::framework::DefaultGradOpMaker<paddle::imperative::OpBase, true>);
-REGISTER_OPERATOR(pool2d_grad, ops::PoolOpGrad);
+REGISTER_OPERATOR(pool2d_grad, ops::PoolOpGrad,
+                  ops::Pool2dOpGradGradMaker<paddle::framework::OpDesc>,
+                  ops::Pool2dOpGradGradMaker<paddle::imperative::OpBase>);
+REGISTER_OPERATOR(pool2d_grad_grad, ops::PoolOp);
 
 REGISTER_OP_CPU_KERNEL(
     pool2d, ops::PoolKernel<paddle::platform::CPUDeviceContext, float>,
@@ -695,6 +712,10 @@ REGISTER_OP_CPU_KERNEL(
 REGISTER_OP_CPU_KERNEL(
     pool2d_grad, ops::PoolGradKernel<paddle::platform::CPUDeviceContext, float>,
     ops::PoolGradKernel<paddle::platform::CPUDeviceContext, double>);
+REGISTER_OP_CPU_KERNEL(
+    pool2d_grad_grad,
+    ops::PoolGradGradKernel<paddle::platform::CPUDeviceContext, float>,
+    ops::PoolGradGradKernel<paddle::platform::CPUDeviceContext, double>);
 
 REGISTER_OPERATOR(
     pool3d, ops::PoolOp, ops::Pool3dOpMaker, ops::PoolOpInferVarType,

paddle/fluid/operators/pool_op.cu

@@ -28,6 +28,13 @@ REGISTER_OP_CUDA_KERNEL(
     ops::PoolGradKernel<paddle::platform::CUDADeviceContext,
                         paddle::platform::float16>);
 
+REGISTER_OP_CUDA_KERNEL(
+    pool2d_grad_grad,
+    ops::PoolGradGradKernel<paddle::platform::CUDADeviceContext, float>,
+    ops::PoolGradGradKernel<paddle::platform::CUDADeviceContext, double>,
+    ops::PoolGradGradKernel<paddle::platform::CUDADeviceContext,
+                            paddle::platform::float16>);
+
 REGISTER_OP_CUDA_KERNEL(
     pool3d, ops::PoolKernel<paddle::platform::CUDADeviceContext, float>,
     ops::PoolKernel<paddle::platform::CUDADeviceContext, double>,

paddle/fluid/operators/pool_op.h

@@ -357,5 +357,19 @@ class PoolGradKernel : public framework::OpKernel<T> {
   }
 };
 
+template <typename DeviceContext, typename T>
+class PoolGradGradKernel : public PoolKernel<DeviceContext, T> {
+ public:
+  void Compute(const framework::ExecutionContext& context) const override {
+    std::string pooling_type = context.Attr<std::string>("pooling_type");
+    if (pooling_type == "max") {
+      PADDLE_THROW(platform::errors::InvalidArgument(
+          "Pool op grad grad only supports avgpool."));
+    } else {
+      PoolKernel<DeviceContext, T>::Compute(context);
+    }
+  }
+};
+
 }  // namespace operators
 }  // namespace paddle

python/paddle/fluid/tests/unittests/test_nn_grad.py

@@ -381,5 +381,103 @@ class TestConcatDoubleGradCheck(unittest.TestCase):
             self.func(p)
 
 
+class TestAvgPool2DDoubleGradCheckCase1(unittest.TestCase):
+    @prog_scope()
+    def func(self, place):
+        input_NCHW = fluid.layers.data(
+            name="input_NCHW",
+            shape=[2, 3, 5, 5],
+            append_batch_size=False,
+            dtype="float32")
+
+        input_NCHW.persistable = True
+        y = layers.pool2d(input_NCHW, pool_size=2, pool_type="avg")
+        x_arr = np.random.uniform(-1, 1, [2, 3, 5, 5]).astype(np.float32)
+
+        gradient_checker.double_grad_check(
+            [input_NCHW], y, x_init=x_arr, place=place, eps=0.05)
+
+    def test_grad(self):
+        places = [fluid.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(fluid.CUDAPlace(0))
+        for p in places:
+            self.func(p)
+
+
+class TestAvgPool2DDoubleGradCheckCase2(unittest.TestCase):
+    @prog_scope()
+    def func(self, place):
+        input_NHWC = fluid.layers.data(
+            name="input_NHWC",
+            shape=[2, 5, 5, 3],
+            append_batch_size=False,
+            dtype="float32")
+
+        input_NHWC.persistable = True
+        y = layers.pool2d(
+            input_NHWC, pool_size=2, pool_type="avg", data_format="NHWC")
+        x_arr = np.random.uniform(-1, 1, [2, 5, 5, 3]).astype(np.float32)
+
+        gradient_checker.double_grad_check(
+            [input_NHWC], y, x_init=x_arr, place=place, eps=0.05)
+
+    def test_grad(self):
+        places = [fluid.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(fluid.CUDAPlace(0))
+        for p in places:
+            self.func(p)
+
+
+class TestAvgPool2DDoubleGradCheckCase3(unittest.TestCase):
+    @prog_scope()
+    def func(self, place):
+        input_NCHW = fluid.layers.data(
+            name="input_NCHW",
+            shape=[2, 3, 5, 5],
+            append_batch_size=False,
+            dtype="float32")
+
+        input_NCHW.persistable = True
+        y = layers.pool2d(
+            input_NCHW, pool_size=2, pool_type="avg", pool_padding=[1, 1])
+        x_arr = np.random.uniform(-1, 1, [2, 3, 5, 5]).astype(np.float32)
+
+        gradient_checker.double_grad_check(
+            [input_NCHW], y, x_init=x_arr, place=place, eps=0.05)
+
+    def test_grad(self):
+        places = [fluid.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(fluid.CUDAPlace(0))
+        for p in places:
+            self.func(p)
+
+
+class TestAvgPool2DDoubleGradCheckCase4(unittest.TestCase):
+    @prog_scope()
+    def func(self, place):
+        input_NCHW = fluid.layers.data(
+            name="input_NCHW",
+            shape=[2, 3, 5, 5],
+            append_batch_size=False,
+            dtype="float32")
+
+        input_NCHW.persistable = True
+        y = layers.pool2d(input_NCHW, pool_size=[4, 4], pool_type="avg")
+        x_arr = np.random.uniform(-1, 1, [2, 3, 5, 5]).astype(np.float32)
+
+        gradient_checker.double_grad_check(
+            [input_NCHW], y, x_init=x_arr, place=place, eps=0.05)
+
+    def test_grad(self):
+        places = [fluid.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(fluid.CUDAPlace(0))
+        for p in places:
+            self.func(p)
+
+
 if __name__ == "__main__":
     unittest.main()