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

paddle/phi/api/yaml/backward.yaml

@@ -1482,17 +1482,17 @@
     func : matrix_power_grad
 
 - backward_op : max_pool2d_with_index_grad
-  forward : max_pool2d_with_index(Tensor x, int[] kernel_size, int[] strides = {1, 1}, int[] paddings = {0, 0}, bool global_pooling = false, bool adaptive = false) -> Tensor(out), Tensor(mask)
-  args : (Tensor x, Tensor mask, Tensor out_grad, int[] kernel_size, int[] strides, int[] paddings, bool global_pooling, bool adaptive)
+  forward : max_pool2d_with_index(Tensor x, int[] kernel_size, int[] strides = {1, 1}, int[] paddings = {0, 0}, bool global_pooling = false, bool adaptive = false, bool fractional = false, float random_u = 0.0) -> Tensor(out), Tensor(mask)
+  args : (Tensor x, Tensor mask, Tensor out_grad, int[] kernel_size, int[] strides, int[] paddings, bool global_pooling, bool adaptive, bool fractional, float random_u)
   output : Tensor(x_grad)
   infer_meta :
     func : MaxPoolWithIndexGradInferMeta
   kernel :
     func : max_pool2d_with_index_grad
 
 - backward_op : max_pool3d_with_index_grad
-  forward : max_pool3d_with_index(Tensor x, int[] kernel_size, int[] strides = {1, 1, 1}, int[] paddings = {0, 0, 0}, bool global_pooling = false, bool adaptive = false) -> Tensor(out), Tensor(mask)
-  args : (Tensor x, Tensor mask, Tensor out_grad, int[] kernel_size, int[] strides, int[] paddings, bool global_pooling, bool adaptive)
+  forward : max_pool3d_with_index(Tensor x, int[] kernel_size, int[] strides = {1, 1, 1}, int[] paddings = {0, 0, 0}, bool global_pooling = false, bool adaptive = false, bool fractional = false, float random_u = 0.0) -> Tensor(out), Tensor(mask)
+  args : (Tensor x, Tensor mask, Tensor out_grad, int[] kernel_size, int[] strides, int[] paddings, bool global_pooling, bool adaptive, bool fractional, float random_u)
   output : Tensor(x_grad)
   infer_meta :
     func : MaxPoolWithIndexGradInferMeta

paddle/phi/api/yaml/ops.yaml

@@ -1682,7 +1682,7 @@
   backward : matrix_power_grad
 
 - op : max_pool2d_with_index
-  args : (Tensor x, int[] kernel_size, int[] strides= {1, 1}, int[] paddings = {0, 0}, bool global_pooling = false, bool adaptive = false)
+  args : (Tensor x, int[] kernel_size, int[] strides= {1, 1}, int[] paddings = {0, 0}, bool global_pooling = false, bool adaptive = false, bool fractional = false, float random_u = 0.0)
   output : Tensor(out), Tensor(mask)
   infer_meta :
     func : MaxPoolWithIndexInferMeta
@@ -1691,7 +1691,7 @@
   backward : max_pool2d_with_index_grad
 
 - op : max_pool3d_with_index
-  args : (Tensor x, int[] kernel_size, int[] strides = {1, 1, 1}, int[] paddings = {0, 0, 0}, bool global_pooling = false, bool adaptive = false)
+  args : (Tensor x, int[] kernel_size, int[] strides = {1, 1, 1}, int[] paddings = {0, 0, 0}, bool global_pooling = false, bool adaptive = false, bool fractional = false, float random_u = 0.0)
   output : Tensor(out), Tensor(mask)
   infer_meta :
     func : MaxPoolWithIndexInferMeta

paddle/phi/infermeta/backward.cc

@@ -658,6 +658,8 @@ void MaxPoolWithIndexGradInferMeta(const MetaTensor& x,
                                    const std::vector<int>& paddings,
                                    bool global_pooling,
                                    bool adaptive,
+                                   bool fractional,
+                                   float random_u,
                                    MetaTensor* dx) {
   dx->share_meta(x);
 }

paddle/phi/infermeta/backward.h

@@ -310,6 +310,8 @@ void MaxPoolWithIndexGradInferMeta(const MetaTensor& x,
                                    const std::vector<int>& paddings,
                                    bool global_pooling,
                                    bool adaptive,
+                                   bool fractional,
+                                   float random_u,
                                    MetaTensor* dx);
 
 void MeshgridGradInferMeta(const std::vector<const MetaTensor*>& inputs,

paddle/phi/infermeta/unary.cc

@@ -2263,6 +2263,8 @@ void MaxPoolWithIndexInferMeta(const MetaTensor& x,
                                const std::vector<int>& paddings,
                                bool global_pooling,
                                bool adaptive,
+                               bool fractional,
+                               float random_u,
                                MetaTensor* out,
                                MetaTensor* mask,
                                MetaConfig config) {
@@ -2309,7 +2311,7 @@ void MaxPoolWithIndexInferMeta(const MetaTensor& x,
           kernel_size_.size()));
 
   std::vector<int64_t> output_shape({x_dims[0], x_dims[1]});
-  if (adaptive) {
+  if (adaptive || fractional) {
     output_shape.insert(
         output_shape.end(), kernel_size_.begin(), kernel_size_.end());
   } else {

paddle/phi/infermeta/unary.h

@@ -344,6 +344,8 @@ void MaxPoolWithIndexInferMeta(const MetaTensor& x,
                                const std::vector<int>& paddings,
                                bool global_pooling,
                                bool adaptive,
+                               bool fractional,
+                               float random_u,
                                MetaTensor* out,
                                MetaTensor* mask,
                                MetaConfig config = MetaConfig());

paddle/phi/kernels/funcs/pooling.cc

@@ -1571,6 +1571,8 @@ class MaxPool2dWithIndexFunctor<CPUContext, T1, T2> {
                   const std::vector<int>& strides,
                   const std::vector<int>& paddings,
                   bool adaptive,
+                  bool fractional,
+                  float random_u,
                   DenseTensor* output,
                   DenseTensor* mask) {
     const int batch_size = static_cast<int>(input.dims()[0]);
@@ -1592,6 +1594,26 @@ class MaxPool2dWithIndexFunctor<CPUContext, T1, T2> {
     T1* output_data = context.template Alloc<T1>(output);
     T2* mask_data = context.template Alloc<T2>(mask);
 
+    float alpha_height = 0, alpha_width = 0;
+    float u_height = 0, u_width = 0;
+    if (fractional) {
+      float u = 0;
+      if (random_u == 0) {
+        std::uniform_real_distribution<float> dist(0, 1);
+        auto engine = phi::GetCPURandomEngine(0);
+        u = dist(*engine);
+      } else {
+        u = random_u;
+      }
+
+      alpha_height = static_cast<float>(input_height) / output_height;
+      alpha_width = static_cast<float>(input_width) / output_width;
+
+      u_height =
+          FractionalRationalU(u, alpha_height, input_height, output_height);
+      u_width = FractionalRationalU(u, alpha_width, input_width, output_width);
+    }
+
     int hstart = 0, hend = 0;
     int wstart = 0, wend = 0;
     for (int i = 0; i < batch_size; i++) {
@@ -1600,6 +1622,11 @@ class MaxPool2dWithIndexFunctor<CPUContext, T1, T2> {
           if (adaptive) {
             hstart = AdaptStartIndex(ph, input_height, output_height);
             hend = AdaptEndIndex(ph, input_height, output_height);
+          } else if (fractional) {
+            hstart = FractionalStartIndex(ph, alpha_height, u_height);
+            hend = FractionalEndIndex(ph, alpha_height, u_height);
+            hstart = std::max(hstart, 0);
+            hend = std::min(hend, input_height);
           } else {
             hstart = ph * stride_height - padding_height;
             hend = std::min(hstart + ksize_height, input_height);
@@ -1609,6 +1636,11 @@ class MaxPool2dWithIndexFunctor<CPUContext, T1, T2> {
             if (adaptive) {
               wstart = AdaptStartIndex(pw, input_width, output_width);
               wend = AdaptEndIndex(pw, input_width, output_width);
+            } else if (fractional) {
+              wstart = FractionalStartIndex(pw, alpha_width, u_width);
+              wend = FractionalEndIndex(pw, alpha_width, u_width);
+              wstart = std::max(wstart, 0);
+              wend = std::min(wend, input_width);
             } else {
               wstart = pw * stride_width - padding_width;
               wend = std::min(wstart + ksize_width, input_width);
@@ -1653,6 +1685,8 @@ class MaxPool2dWithIndexGradFunctor<CPUContext, T1, T2> {
                   const std::vector<int>& strides UNUSED,
                   const std::vector<int>& paddings UNUSED,
                   bool adaptive UNUSED,
+                  bool fractional UNUSED,
+                  float random_u UNUSED,
                   DenseTensor* input_grad) {
     const int batch_size = static_cast<int>(input_grad->dims()[0]);
     const int input_height = static_cast<int>(input_grad->dims()[2]);
@@ -1704,6 +1738,8 @@ class MaxPool3dWithIndexFunctor<CPUContext, T1, T2> {
                   const std::vector<int>& strides,
                   const std::vector<int>& paddings,
                   bool adaptive,
+                  bool fractional,
+                  float random_u,
                   DenseTensor* output,
                   DenseTensor* mask) {
     const int batch_size = static_cast<int>(input.dims()[0]);
@@ -1730,6 +1766,28 @@ class MaxPool3dWithIndexFunctor<CPUContext, T1, T2> {
     T1* output_data = context.template Alloc<T1>(output);
     T2* mask_data = context.template Alloc<T2>(mask);
 
+    float alpha_height = 0, alpha_width = 0, alpha_depth = 0;
+    float u_height = 0, u_width = 0, u_depth = 0;
+    if (fractional) {
+      float u = 0;
+      if (random_u == 0) {
+        std::uniform_real_distribution<float> dist(0, 1);
+        auto engine = phi::GetCPURandomEngine(0);
+        u = dist(*engine);
+      } else {
+        u = random_u;
+      }
+
+      alpha_depth = static_cast<float>(input_depth) / output_depth;
+      alpha_height = static_cast<float>(input_height) / output_height;
+      alpha_width = static_cast<float>(input_width) / output_width;
+
+      u_depth = FractionalRationalU(u, alpha_depth, input_depth, output_depth);
+      u_height =
+          FractionalRationalU(u, alpha_height, input_height, output_height);
+      u_width = FractionalRationalU(u, alpha_width, input_width, output_width);
+    }
+
     int dstart = 0, dend = 0;
     int hstart = 0, hend = 0;
     int wstart = 0, wend = 0;
@@ -1739,6 +1797,11 @@ class MaxPool3dWithIndexFunctor<CPUContext, T1, T2> {
           if (adaptive) {
             dstart = AdaptStartIndex(pd, input_depth, output_depth);
             dend = AdaptEndIndex(pd, input_depth, output_depth);
+          } else if (fractional) {
+            dstart = FractionalStartIndex(pd, alpha_depth, u_depth);
+            dend = FractionalEndIndex(pd, alpha_depth, u_depth);
+            dstart = std::max(dstart, 0);
+            dend = std::min(dend, input_depth);
           } else {
             dstart = pd * stride_depth - padding_depth;
             dend = std::min(dstart + ksize_depth, input_depth);
@@ -1748,6 +1811,11 @@ class MaxPool3dWithIndexFunctor<CPUContext, T1, T2> {
             if (adaptive) {
               hstart = AdaptStartIndex(ph, input_height, output_height);
               hend = AdaptEndIndex(ph, input_height, output_height);
+            } else if (fractional) {
+              hstart = FractionalStartIndex(ph, alpha_height, u_height);
+              hend = FractionalEndIndex(ph, alpha_height, u_height);
+              hstart = std::max(hstart, 0);
+              hend = std::min(hend, input_height);
             } else {
               hstart = ph * stride_height - padding_height;
               hend = std::min(hstart + ksize_height, input_height);
@@ -1757,6 +1825,11 @@ class MaxPool3dWithIndexFunctor<CPUContext, T1, T2> {
               if (adaptive) {
                 wstart = AdaptStartIndex(pw, input_width, output_width);
                 wend = AdaptEndIndex(pw, input_width, output_width);
+              } else if (fractional) {
+                wstart = FractionalStartIndex(pw, alpha_width, u_width);
+                wend = FractionalEndIndex(pw, alpha_width, u_width);
+                wstart = std::max(wstart, 0);
+                wend = std::min(wend, input_width);
               } else {
                 wstart = pw * stride_width - padding_width;
                 wend = std::min(wstart + ksize_width, input_width);
@@ -1806,6 +1879,8 @@ class MaxPool3dWithIndexGradFunctor<CPUContext, T1, T2> {
                   const std::vector<int>& strides UNUSED,
                   const std::vector<int>& paddings UNUSED,
                   bool adaptive UNUSED,
+                  bool fractional UNUSED,
+                  float random_u UNUSED,
                   DenseTensor* input_grad) {
     const int batch_size = static_cast<int>(input_grad->dims()[0]);
     const int input_depth = static_cast<int>(input_grad->dims()[2]);