【Hackathon 6th No.29】为 paddle.nn.functional.max_unpool1d / max_unpool2d / max_unpool3d 进行功能增强 -part (#63648)

NKNaN · web-flow · commit 08c0424a5951 · 2024-04-22T12:54:12.000+08:00
diff --git a/paddle/phi/kernels/cpu/unpool_grad_kernel.cc b/paddle/phi/kernels/cpu/unpool_grad_kernel.cc
@@ -130,8 +130,18 @@ void Unpool3dGradKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    unpool_grad, CPU, ALL_LAYOUT, phi::UnpoolGradKernel, float, double) {}
+PD_REGISTER_KERNEL(unpool_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::UnpoolGradKernel,
+                   float,
+                   double,
+                   int64_t) {}
 
-PD_REGISTER_KERNEL(
-    unpool3d_grad, CPU, ALL_LAYOUT, phi::Unpool3dGradKernel, float, double) {}
+PD_REGISTER_KERNEL(unpool3d_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::Unpool3dGradKernel,
+                   float,
+                   double,
+                   int64_t) {}
diff --git a/paddle/phi/kernels/cpu/unpool_kernel.cc b/paddle/phi/kernels/cpu/unpool_kernel.cc
@@ -126,7 +126,8 @@ void Unpool3dKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(unpool, CPU, ALL_LAYOUT, phi::UnpoolKernel, float, double) {}
+PD_REGISTER_KERNEL(
+    unpool, CPU, ALL_LAYOUT, phi::UnpoolKernel, float, double, int64_t) {}
 
 PD_REGISTER_KERNEL(
-    unpool3d, CPU, ALL_LAYOUT, phi::Unpool3dKernel, float, double) {}
+    unpool3d, CPU, ALL_LAYOUT, phi::Unpool3dKernel, float, double, int64_t) {}
diff --git a/paddle/phi/kernels/gpu/unpool_grad_kernel.cu b/paddle/phi/kernels/gpu/unpool_grad_kernel.cu
@@ -188,8 +188,18 @@ void Unpool3dGradKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    unpool_grad, GPU, ALL_LAYOUT, phi::UnpoolGradKernel, float, double) {}
+PD_REGISTER_KERNEL(unpool_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::UnpoolGradKernel,
+                   float,
+                   double,
+                   int64_t) {}
 
-PD_REGISTER_KERNEL(
-    unpool3d_grad, GPU, ALL_LAYOUT, phi::Unpool3dGradKernel, float, double) {}
+PD_REGISTER_KERNEL(unpool3d_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::Unpool3dGradKernel,
+                   float,
+                   double,
+                   int64_t) {}
diff --git a/paddle/phi/kernels/gpu/unpool_kernel.cu b/paddle/phi/kernels/gpu/unpool_kernel.cu
@@ -173,7 +173,13 @@ void Unpool3dKernel(const Context& dev_ctx,
 }  // namespace phi
 
 PD_REGISTER_KERNEL(
-    unpool, GPU, ALL_LAYOUT, phi::UnpoolKernel, int, float, double) {}
+    unpool, GPU, ALL_LAYOUT, phi::UnpoolKernel, int, float, double, int64_t) {}
 
-PD_REGISTER_KERNEL(
-    unpool3d, GPU, ALL_LAYOUT, phi::Unpool3dKernel, int, float, double) {}
+PD_REGISTER_KERNEL(unpool3d,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::Unpool3dKernel,
+                   int,
+                   float,
+                   double,
+                   int64_t) {}
diff --git a/python/paddle/nn/functional/pooling.py b/python/paddle/nn/functional/pooling.py
@@ -755,7 +755,7 @@ def max_unpool1d(
         x (Tensor): The input tensor of unpooling operator which is a 3-D tensor with
                           shape [N, C, L]. The format of input tensor is `"NCL"`,
                           where `N` is batch size, `C` is the number of channels, `L` is
-                          the length of the feature. The data type is float32 or float64.
+                          the length of the feature. The data type is float32, float64 or int64.
         indices (Tensor): The indices given out by maxpooling1d which is a 3-D tensor with
                           shape [N, C, L]. The format of input tensor is `"NCL"` ,
                           where `N` is batch size, `C` is the number of channels, `L` is
@@ -813,6 +813,8 @@ def max_unpool1d(
     # use 2d to implenment 1d should expand padding in advance.
     padding = _expand_low_nd_padding(padding)
 
+    if output_size is not None:
+        output_size = output_size[:2] + [1] + output_size[2:]
     output_size = _unpool_output_size(
         x, kernel_size, stride, padding, output_size
     )
@@ -863,12 +865,12 @@ def max_unpool2d(
                           shape [N, C, H, W]. The format of input tensor is `"NCHW"`,
                           where `N` is batch size, `C` is the number of channels,
                           `H` is the height of the feature, and `W` is the width of the
-                          feature. The data type if float32 or float64.
+                          feature. The data type is float32, float64 or int64.
         indices (Tensor): The indices given out by maxpooling2d which is a 4-D tensor with
                           shape [N, C, H, W]. The format of input tensor is `"NCHW"` ,
                           where `N` is batch size, `C` is the number of channels,
                           `H` is the height of the feature, and `W` is the width of the
-                          feature. The data type if float32 or float64.
+                          feature. The data type is float32 or float64.
         kernel_size (int|list|tuple): The unpool kernel size. If unpool kernel size is a tuple or list,
             it must contain an integer.
         stride (int|list|tuple): The unpool stride size. If unpool stride size is a tuple or list,
@@ -1011,7 +1013,7 @@ def max_unpool3d(
                           shape [N, C, D, H, W]. The format of input tensor is `"NCDHW"`,
                           where `N` is batch size, `C` is the number of channels, `D` is
                           the depth of the feature, `H` is the height of the feature,
-                          and `W` is the width of the feature. The data type is float32 or float64.
+                          and `W` is the width of the feature. The data type is float32, float64 or int64.
         indices (Tensor): The indices given out by maxpooling3d which is a 5-D tensor with
                           shape [N, C, D, H, W]. The format of input tensor is `"NCDHW"` ,
                           where `N` is batch size, `C` is the number of channels, `D` is
diff --git a/test/deprecated/legacy_test/test_unpool3d_op.py b/test/deprecated/legacy_test/test_unpool3d_op.py
@@ -373,6 +373,44 @@ def test_case(self):
         paddle.enable_static()
 
 
+class TestUnpool3DOpAPI_dygraph4(unittest.TestCase):
+    def test_case(self):
+        places = [paddle.CPUPlace()]
+        if paddle.base.core.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        for place in places:
+            paddle.disable_static()
+            input_data = (
+                np.arange(3 * 4 * 4 * 6)
+                .reshape([1, 3, 4, 4, 6])
+                .astype("float32")
+            )
+            input_x = paddle.to_tensor(input_data)
+            output, indices = F.max_pool3d(
+                input_x, kernel_size=2, stride=2, return_mask=True
+            )
+            output_unpool = F.max_unpool3d(
+                output.astype("int64"),
+                indices,
+                kernel_size=2,
+                stride=2,
+                output_size=input_x.shape,
+            )
+            expected_output_unpool = unpool3dmax_forward_naive(
+                output.numpy(),
+                indices.numpy(),
+                [2, 2, 2],
+                [2, 2, 2],
+                [0, 0, 0],
+                [4, 4, 6],
+            )
+            np.testing.assert_allclose(
+                output_unpool.numpy(), expected_output_unpool, rtol=1e-05
+            )
+
+        paddle.enable_static()
+
+
 class TestUnpool3DOpAPI_static(unittest.TestCase):
     @test_with_pir_api
     def test_case(self):
diff --git a/test/deprecated/legacy_test/test_unpool_op.py b/test/deprecated/legacy_test/test_unpool_op.py
@@ -400,6 +400,51 @@ def test_case(self):
             np.testing.assert_allclose(out_pp.numpy(), expect_res, rtol=1e-05)
 
 
+class TestUnpoolOpAPI_dy4(unittest.TestCase):
+    def test_case(self):
+        import numpy as np
+
+        import paddle
+        import paddle.nn.functional as F
+        from paddle import base
+        from paddle.base import core
+
+        if core.is_compiled_with_cuda():
+            place = core.CUDAPlace(0)
+        else:
+            place = core.CPUPlace()
+        with base.dygraph.guard(place):
+            input_data = np.array(
+                [
+                    [
+                        [
+                            [1, 2, 3, 4, 5],
+                            [6, 7, 8, 9, 10],
+                            [11, 12, 13, 14, 15],
+                            [16, 17, 18, 19, 20],
+                        ]
+                    ]
+                ]
+            ).astype("float32")
+            input_x = paddle.to_tensor(input_data)
+            output, indices = F.max_pool2d(
+                input_x, kernel_size=2, stride=2, return_mask=True
+            )
+            out_pp = F.max_unpool2d(
+                output.astype("int64"),
+                indices,
+                kernel_size=2,
+                stride=None,
+                output_size=input_x.shape,
+            )
+            output_np = output.numpy()
+            indices_np = indices.numpy()
+            expect_res = unpool2dmax_forward_naive(
+                output_np, indices_np, [2, 2], [2, 2], [0, 0], [4, 5]
+            ).astype("float64")
+            np.testing.assert_allclose(out_pp.numpy(), expect_res, rtol=1e-05)
+
+
 class TestUnpoolOpAPI_st(unittest.TestCase):
     @test_with_pir_api
     def test_case(self):
diff --git a/test/legacy_test/test_unpool1d_op.py b/test/legacy_test/test_unpool1d_op.py
@@ -135,6 +135,35 @@ def test_case(self):
         paddle.enable_static()
 
 
+class TestUnpool1DOpAPI_dygraph4(unittest.TestCase):
+    def test_case(self):
+        places = [paddle.CPUPlace()]
+        if paddle.base.core.is_compiled_with_cuda():
+            places.append(paddle.CUDAPlace(0))
+        for place in places:
+            paddle.disable_static()
+            input_data = np.arange(3 * 16).reshape([1, 3, 16]).astype("float32")
+            input_x = paddle.to_tensor(input_data)
+            output, indices = F.max_pool1d(
+                input_x, kernel_size=2, stride=2, return_mask=True
+            )
+            output_unpool = F.max_unpool1d(
+                output.astype("int64"),
+                indices,
+                kernel_size=2,
+                stride=2,
+                output_size=input_x.shape,
+            )
+            expected_output_unpool = unpool1dmax_forward_naive(
+                output.numpy(), indices.numpy(), [2], [2], [0], [16]
+            )
+            np.testing.assert_allclose(
+                output_unpool.numpy(), expected_output_unpool, rtol=1e-05
+            )
+
+        paddle.enable_static()
+
+
 class TestUnpool1DOpAPI_static(unittest.TestCase):
     @test_with_pir_api
     def test_case(self):
