F.pad data_format support automatic default format 易用性提升 (#68214)

NKNaN · web-flow · commit 5d2ff7352d6a · 2024-09-23T19:16:56.000+08:00
* F.pad data_format support auto infer

* fix typing

* fix code style
diff --git a/python/paddle/nn/functional/common.py b/python/paddle/nn/functional/common.py
@@ -1690,7 +1690,7 @@ def pad(
     pad: ShapeLike,
     mode: _PaddingTensorMode = 'constant',
     value: float = 0.0,
-    data_format: DataLayoutND = "NCHW",
+    data_format: DataLayoutND | None = None,
     pad_from_left_axis: bool = True,
     name: str | None = None,
 ) -> Tensor:
@@ -1733,7 +1733,9 @@ def pad(
         value (float, optional): The value to fill the padded areas in 'constant' mode . Default is :math:`0.0`.
         data_format (str, optional): An string from: ``'NCL'``, ``'NLC'``, ``'NHWC'``, ``'NCHW'``, ``'NCDHW'``, ``'NDHWC'``. Specify the data format of
            the input data when: 1. mode is any of ``'reflect'``, ``'replicate'`` or ``'circular'``; or 2. the input ``'pad'`` is a tensor;
-           or 3. the length of ``'pad'`` is ``2*(x.ndim - 2)``. Default: ``'NCHW'``.
+           or 3. the length of ``'pad'`` is ``2*(x.ndim - 2)``. The default value is None, which means it will be automatically inferred from the
+           input dimension of ``'x'``. When ``'x'`` is a 3-D Tensor, data_format will be set to ``'NCL'``; When ``'x'`` is a 4-D Tensor,
+           data_format will be set to ``'NCHW'``; When ``'x'`` is a 5-D Tensor, data_format will be set to ``'NCDHW'``.
         pad_from_left_axis (bool, optional): The parameter is only valid when mode is ``'constant'`` and the input ``'pad'`` is
            length of ``'pad'`` is ``2*x.ndim``, the order of padding can be customized. If True, the padding will be started from
            the first axis of ``'x'``; if False, it will be started from the last axis of ``'x'``. Default: True.
@@ -1888,12 +1890,6 @@ def pad(
         'circular',
     ], f"mode should be one of constant, reflect, replicate, circular, but got {mode}."
 
-    data_format = data_format.upper()
-    assert data_format in ["NCL", "NCHW", "NCDHW", "NLC", "NHWC", "NDHWC"], (
-        "data_format should be in one of [NCL, NCHW, NCDHW, NLC, NHWC, NDHWC], "
-        f"but got {data_format}"
-    )
-
     x_dim = len(x.shape)
 
     if (
@@ -1966,6 +1962,19 @@ def pad(
         5,
     ], f"input tensor dimension must be in [3, 4, 5] but got {x_dim}"
 
+    if data_format is None:
+        if x_dim == 3:
+            data_format = "NCL"
+        elif x_dim == 4:
+            data_format = "NCHW"
+        elif x_dim == 5:
+            data_format = "NCDHW"
+
+    data_format = data_format.upper()
+    assert data_format in ["NCL", "NCHW", "NCDHW", "NLC", "NHWC", "NDHWC"], (
+        "data_format should be in one of [NCL, NCHW, NCDHW, NLC, NHWC, NDHWC], "
+        f"but got {data_format}"
+    )
     supported_format_map = {
         3: ["NCL", "NLC"],
         4: ["NCHW", "NHWC"],
diff --git a/test/legacy_test/test_pad3d_op.py b/test/legacy_test/test_pad3d_op.py
@@ -447,11 +447,12 @@ def check_static_result_2(self, place):
             )
             result1 = F.pad(x=x, pad=pad, mode=mode, data_format="NCDHW")
             result2 = F.pad(x=x, pad=pad, mode=mode, data_format="NDHWC")
+            result3 = F.pad(x=x, pad=pad, mode=mode)
             exe = Executor(place)
             fetches = exe.run(
                 paddle.static.default_main_program(),
                 feed={"x": input_data},
-                fetch_list=[result1, result2],
+                fetch_list=[result1, result2, result3],
             )
 
             np_out1 = self._get_numpy_out(
@@ -462,6 +463,7 @@ def check_static_result_2(self, place):
             )
             np.testing.assert_allclose(fetches[0], np_out1, rtol=1e-05)
             np.testing.assert_allclose(fetches[1], np_out2, rtol=1e-05)
+            np.testing.assert_allclose(fetches[2], np_out1, rtol=1e-05)
 
     def check_static_result_3(self, place):
         paddle.enable_static()
@@ -482,11 +484,12 @@ def check_static_result_3(self, place):
             )
             result1 = F.pad(x=x, pad=pad, mode=mode, data_format="NCDHW")
             result2 = F.pad(x=x, pad=pad, mode=mode, data_format="NDHWC")
+            result3 = F.pad(x=x, pad=pad, mode=mode)
             exe = Executor(place)
             fetches = exe.run(
                 paddle.static.default_main_program(),
                 feed={"x": input_data},
-                fetch_list=[result1, result2],
+                fetch_list=[result1, result2, result3],
             )
 
             np_out1 = self._get_numpy_out(
@@ -497,6 +500,7 @@ def check_static_result_3(self, place):
             )
             np.testing.assert_allclose(fetches[0], np_out1, rtol=1e-05)
             np.testing.assert_allclose(fetches[1], np_out2, rtol=1e-05)
+            np.testing.assert_allclose(fetches[2], np_out1, rtol=1e-05)
 
     def check_static_result_4(self, place):
         paddle.enable_static()
@@ -517,11 +521,12 @@ def check_static_result_4(self, place):
             )
             result1 = F.pad(x=x, pad=pad, mode=mode, data_format="NCDHW")
             result2 = F.pad(x=x, pad=pad, mode=mode, data_format="NDHWC")
+            result3 = F.pad(x=x, pad=pad, mode=mode)
             exe = Executor(place)
             fetches = exe.run(
                 paddle.static.default_main_program(),
                 feed={"x": input_data},
-                fetch_list=[result1, result2],
+                fetch_list=[result1, result2, result3],
             )
 
             np_out1 = self._get_numpy_out(
@@ -532,6 +537,7 @@ def check_static_result_4(self, place):
             )
             np.testing.assert_allclose(fetches[0], np_out1, rtol=1e-05)
             np.testing.assert_allclose(fetches[1], np_out2, rtol=1e-05)
+            np.testing.assert_allclose(fetches[2], np_out1, rtol=1e-05)
 
     def _get_numpy_out(
         self, input_data, pad, mode, value=0, data_format="NCDHW"
@@ -632,10 +638,12 @@ def test_dygraph_1(self):
         y3 = F.pad(
             tensor_data, pad=pad_3, mode=mode, value=value, data_format="NCDHW"
         )
+        y4 = F.pad(tensor_data, pad=pad, mode=mode, value=value)
 
         np.testing.assert_allclose(y1.numpy(), np_out1, rtol=1e-05)
         np.testing.assert_allclose(y2.numpy(), np_out2, rtol=1e-05)
         np.testing.assert_allclose(y3.numpy(), np_out3, rtol=1e-05)
+        np.testing.assert_allclose(y4.numpy(), np_out1, rtol=1e-05)
 
     def test_dygraph_2(self):
         paddle.disable_static()
@@ -675,10 +683,16 @@ def test_dygraph_2(self):
         y3 = F.pad(
             tensor_data, pad=pad_3, mode=mode, value=value, data_format="NCHW"
         )
-
+        y4 = F.pad(
+            tensor_data,
+            pad=tensor_pad,
+            mode=mode,
+            value=value,
+        )
         np.testing.assert_allclose(y1.numpy(), np_out1, rtol=1e-05)
         np.testing.assert_allclose(y2.numpy(), np_out2, rtol=1e-05)
         np.testing.assert_allclose(y3.numpy(), np_out3, rtol=1e-05)
+        np.testing.assert_allclose(y4.numpy(), np_out1, rtol=1e-05)
 
     def test_dygraph_3(self):
         paddle.disable_static()
@@ -721,10 +735,17 @@ def test_dygraph_3(self):
         y3 = F.pad(
             tensor_data, pad=pad_3, mode=mode, value=value, data_format="NCL"
         )
+        y4 = F.pad(
+            tensor_data,
+            pad=tensor_pad,
+            mode=mode,
+            value=value,
+        )
 
         np.testing.assert_allclose(y1.numpy(), np_out1, rtol=1e-05)
         np.testing.assert_allclose(y2.numpy(), np_out2, rtol=1e-05)
         np.testing.assert_allclose(y3.numpy(), np_out3, rtol=1e-05)
+        np.testing.assert_allclose(y4.numpy(), np_out1, rtol=1e-05)
 
 
 class TestPadAPI_complex64(TestPadAPI):
