【Hackathon 5th No.38】为 Paddle 新增 FractionalMaxPool2d / FractionalMaxPool3d API -kernel

paddle/phi/api/yaml/backward.yaml

@@ -911,6 +911,26 @@
     data_type : out_grad
   no_need_buffer : x
 
+- backward_op : fractional_max_pool2d_with_index_grad
+  forward : fractional_max_pool2d_with_index(Tensor x, int[] output_size, float random_u = 0.0) -> Tensor(out), Tensor(mask)
+  args : (Tensor x, Tensor mask, Tensor out_grad, int[] output_size, float random_u)
+  output : Tensor(x_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param : [x]
+  kernel :
+    func : fractional_max_pool2d_with_index_grad
+
+- backward_op : fractional_max_pool3d_with_index_grad
+  forward : fractional_max_pool3d_with_index(Tensor x, int[] output_size, float random_u = 0.0) -> Tensor(out), Tensor(mask)
+  args : (Tensor x, Tensor mask, Tensor out_grad, int[] output_size, float random_u)
+  output : Tensor(x_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param : [x]
+  kernel :
+    func : fractional_max_pool3d_with_index_grad
+
 - backward_op : frame_grad
   forward : frame(Tensor x, int frame_length, int hop_length, int axis=-1) -> Tensor(out)
   args : (Tensor x, Tensor out_grad, int frame_length, int hop_length, int axis)

paddle/phi/api/yaml/op_compat.yaml

@@ -1158,6 +1158,18 @@
   outputs :
     out : Y
 
+- op : fractional_max_pool2d_with_index
+  inputs :
+    {x : X}
+  outputs :
+    {out : Out, mask : Mask}
+
+- op : fractional_max_pool3d_with_index
+  inputs :
+    {x : X}
+  outputs :
+    {out : Out, mask : Mask}
+
 - op : frame
   backward : frame_grad
   inputs :

paddle/phi/api/yaml/ops.yaml

@@ -1019,6 +1019,24 @@
     func: fold
   backward: fold_grad
 
+- op : fractional_max_pool2d_with_index
+  args : (Tensor x, int[] output_size, float random_u = 0.0)
+  output : Tensor(out), Tensor(mask)
+  infer_meta :
+    func : FractionalMaxPoolWithIndexInferMeta
+  kernel :
+    func : fractional_max_pool2d_with_index
+  backward : fractional_max_pool2d_with_index_grad
+
+- op : fractional_max_pool3d_with_index
+  args : (Tensor x, int[] output_size, float random_u = 0.0)
+  output : Tensor(out), Tensor(mask)
+  infer_meta :
+    func : FractionalMaxPoolWithIndexInferMeta
+  kernel :
+    func : fractional_max_pool3d_with_index
+  backward : fractional_max_pool3d_with_index_grad
+
 - op : frame
   args : (Tensor x, int frame_length, int hop_length, int axis=-1)
   output : Tensor(out)

paddle/phi/infermeta/unary.cc

@@ -2343,6 +2343,42 @@ void MaxPoolWithIndexInferMeta(const MetaTensor& x,
   mask->set_dtype(phi::CppTypeToDataType<int>::Type());
 }
 
+void FractionalMaxPoolWithIndexInferMeta(const MetaTensor& x,
+                                         const std::vector<int>& output_size,
+                                         float random_u,
+                                         MetaTensor* out,
+                                         MetaTensor* mask,
+                                         MetaConfig config) {
+  std::vector<int> output_size_ = output_size;
+
+  auto x_dims = x.dims();
+
+  PADDLE_ENFORCE_EQ(
+      (x_dims.size() == 4 || x_dims.size() == 5),
+      true,
+      errors::InvalidArgument("Pooling intput should be 4-D or "
+                              "5-D tensor but received %dD-Tensor",
+                              x_dims.size()));
+
+  PADDLE_ENFORCE_EQ(
+      x_dims.size() - output_size_.size(),
+      2U,
+      errors::InvalidArgument(
+          "The input size %d minus the output size %d should equal to 2.",
+          x_dims.size(),
+          output_size_.size()));
+
+  std::vector<int64_t> output_shape({x_dims[0], x_dims[1]});
+  output_shape.insert(
+      output_shape.end(), output_size_.begin(), output_size_.end());
+
+  out->set_dims(common::make_ddim(output_shape));
+  out->set_dtype(x.dtype());
+
+  mask->set_dims(common::make_ddim(output_shape));
+  mask->set_dtype(phi::CppTypeToDataType<int>::Type());
+}
+
 void MeanAllInferMeta(const MetaTensor& x, MetaTensor* out) {
   out->set_dims(common::make_ddim({}));
   out->set_dtype(x.dtype());

paddle/phi/infermeta/unary.h

@@ -348,6 +348,13 @@ void MaxPoolWithIndexInferMeta(const MetaTensor& x,
                                MetaTensor* mask,
                                MetaConfig config = MetaConfig());
 
+void FractionalMaxPoolWithIndexInferMeta(const MetaTensor& x,
+                                         const std::vector<int>& output_size,
+                                         float random_u,
+                                         MetaTensor* out,
+                                         MetaTensor* mask,
+                                         MetaConfig config = MetaConfig());
+
 void MeanAllInferMeta(const MetaTensor& x, MetaTensor* out);
 
 void ModeInferMeta(const MetaTensor& x,

paddle/phi/kernels/cpu/pool_grad_kernel.cc

@@ -44,3 +44,23 @@ PD_REGISTER_KERNEL(max_pool3d_with_index_grad,
                    double) {
   kernel->InputAt(1).SetDataType(phi::CppTypeToDataType<int>::Type());
 }
+
+PD_REGISTER_KERNEL(fractional_max_pool2d_with_index_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::FractionalMaxPool2dWithIndexGradKernel,
+                   float,
+                   double,
+                   phi::dtype::float16) {
+  kernel->InputAt(1).SetDataType(phi::CppTypeToDataType<int>::Type());
+}
+
+PD_REGISTER_KERNEL(fractional_max_pool3d_with_index_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::FractionalMaxPool3dWithIndexGradKernel,
+                   float,
+                   double,
+                   phi::dtype::float16) {
+  kernel->InputAt(1).SetDataType(phi::CppTypeToDataType<int>::Type());
+}

paddle/phi/kernels/cpu/pool_kernel.cc

@@ -36,3 +36,23 @@ PD_REGISTER_KERNEL(max_pool3d_with_index,
                    double) {
   kernel->OutputAt(1).SetDataType(phi::CppTypeToDataType<int>::Type());
 }
+
+PD_REGISTER_KERNEL(fractional_max_pool2d_with_index,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::FractionalMaxPool2dWithIndexKernel,
+                   float,
+                   double,
+                   phi::dtype::float16) {
+  kernel->OutputAt(1).SetDataType(phi::CppTypeToDataType<int>::Type());
+}
+
+PD_REGISTER_KERNEL(fractional_max_pool3d_with_index,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::FractionalMaxPool3dWithIndexKernel,
+                   float,
+                   double,
+                   phi::dtype::float16) {
+  kernel->OutputAt(1).SetDataType(phi::CppTypeToDataType<int>::Type());
+}