[PIR] Migrate paddle.nn.MSELoss into pir (#58352)

MarioLulab · web-flow · commit ff17b89b98b4 · 2023-10-25T11:40:37.000+08:00
diff --git a/test/legacy_test/test_mse_loss.py b/test/legacy_test/test_mse_loss.py
@@ -20,39 +20,47 @@
 from paddle import base
 from paddle.base import core
 from paddle.base.executor import Executor
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestMseLoss(unittest.TestCase):
+    @test_with_pir_api
     def test_mse_loss(self):
         input_val = np.random.uniform(0.1, 0.5, (2, 3)).astype("float32")
         label_val = np.random.uniform(0.1, 0.5, (2, 3)).astype("float32")
 
         sub = input_val - label_val
         np_result = np.mean(sub * sub)
 
-        input_var = paddle.static.data(
-            name="input", shape=[-1, 3], dtype="float32"
-        )
-        label_var = paddle.static.data(
-            name="label", shape=[-1, 3], dtype="float32"
-        )
-
-        output = paddle.nn.functional.mse_loss(input=input_var, label=label_var)
-        for use_cuda in (
-            [False, True] if core.is_compiled_with_cuda() else [False]
-        ):
-            place = base.CUDAPlace(0) if use_cuda else base.CPUPlace()
-            exe = Executor(place)
-            (result,) = exe.run(
-                base.default_main_program(),
-                feed={"input": input_val, "label": label_val},
-                fetch_list=[output],
+        main = paddle.static.Program()
+        startup = paddle.static.Program()
+        with paddle.static.program_guard(main, startup):
+            input_var = paddle.static.data(
+                name="input", shape=[-1, 3], dtype="float32"
             )
+            label_var = paddle.static.data(
+                name="label", shape=[-1, 3], dtype="float32"
+            )
+
+            output = paddle.nn.functional.mse_loss(
+                input=input_var, label=label_var
+            )
+            for use_cuda in (
+                [False, True] if core.is_compiled_with_cuda() else [False]
+            ):
+                place = base.CUDAPlace(0) if use_cuda else base.CPUPlace()
+                exe = Executor(place)
+                (result,) = exe.run(
+                    main,
+                    feed={"input": input_val, "label": label_val},
+                    fetch_list=[output],
+                )
 
-            np.testing.assert_allclose(np_result, result, rtol=1e-05)
+                np.testing.assert_allclose(np_result, result, rtol=1e-05)
 
 
 class TestMseInvalidInput(unittest.TestCase):
+    @test_with_pir_api
     def test_error(self):
         def test_invalid_input():
             input = [256, 3]
@@ -74,6 +82,7 @@ def test_invalid_label():
 
 
 class TestNNMseLoss(unittest.TestCase):
+    @test_with_pir_api
     def test_NNMseLoss_mean(self):
         for dim in [[10, 10], [2, 10, 10], [3, 3, 10, 10]]:
             input_np = np.random.uniform(0.1, 0.5, dim).astype("float32")
@@ -88,13 +97,11 @@ def test_NNMseLoss_mean(self):
             )
             with base.program_guard(prog, startup_prog):
                 input = paddle.static.data(
-                    name='input', shape=[-1] + dim, dtype='float32'
+                    name='input', shape=dim, dtype='float32'
                 )
-                input.desc.set_need_check_feed(False)
                 label = paddle.static.data(
-                    name='label', shape=[-1] + dim, dtype='float32'
+                    name='label', shape=dim, dtype='float32'
                 )
-                label.desc.set_need_check_feed(False)
                 mse_loss = paddle.nn.loss.MSELoss()
                 ret = mse_loss(input, label)
 
@@ -120,6 +127,7 @@ def test_NNMseLoss_mean(self):
             np.testing.assert_allclose(dy_result, expected, rtol=1e-05)
             self.assertEqual(dy_result.shape, ())
 
+    @test_with_pir_api
     def test_NNMseLoss_sum(self):
         for dim in [[10, 10], [2, 10, 10], [3, 3, 10, 10]]:
             input_np = np.random.uniform(0.1, 0.5, dim).astype("float32")
@@ -134,13 +142,11 @@ def test_NNMseLoss_sum(self):
             )
             with base.program_guard(prog, startup_prog):
                 input = paddle.static.data(
-                    name='input', shape=[-1] + dim, dtype='float32'
+                    name='input', shape=dim, dtype='float32'
                 )
-                input.desc.set_need_check_feed(False)
                 label = paddle.static.data(
-                    name='label', shape=[-1] + dim, dtype='float32'
+                    name='label', shape=dim, dtype='float32'
                 )
-                label.desc.set_need_check_feed(False)
                 mse_loss = paddle.nn.loss.MSELoss(reduction='sum')
                 ret = mse_loss(input, label)
 
@@ -166,6 +172,7 @@ def test_NNMseLoss_sum(self):
             np.testing.assert_allclose(dy_result, expected, rtol=1e-05)
             self.assertEqual(dy_result.shape, ())
 
+    @test_with_pir_api
     def test_NNMseLoss_none(self):
         for dim in [[10, 10], [2, 10, 10], [3, 3, 10, 10]]:
             input_np = np.random.uniform(0.1, 0.5, dim).astype("float32")
@@ -180,13 +187,11 @@ def test_NNMseLoss_none(self):
             )
             with base.program_guard(prog, startup_prog):
                 input = paddle.static.data(
-                    name='input', shape=[-1] + dim, dtype='float32'
+                    name='input', shape=dim, dtype='float32'
                 )
-                input.desc.set_need_check_feed(False)
                 label = paddle.static.data(
-                    name='label', shape=[-1] + dim, dtype='float32'
+                    name='label', shape=dim, dtype='float32'
                 )
-                label.desc.set_need_check_feed(False)
                 mse_loss = paddle.nn.loss.MSELoss(reduction='none')
                 ret = mse_loss(input, label)
 
@@ -214,6 +219,7 @@ def test_NNMseLoss_none(self):
 
 
 class TestNNFunctionalMseLoss(unittest.TestCase):
+    @test_with_pir_api
     def test_NNFunctionalMseLoss_mean(self):
         for dim in [[10, 10], [2, 10, 10], [3, 3, 10, 10]]:
             input_np = np.random.uniform(0.1, 0.5, dim).astype("float32")
@@ -256,6 +262,7 @@ def test_NNFunctionalMseLoss_mean(self):
             np.testing.assert_allclose(dy_result, expected, rtol=1e-05)
             self.assertEqual(dy_result.shape, ())
 
+    @test_with_pir_api
     def test_NNFunctionalMseLoss_sum(self):
         for dim in [[10, 10], [2, 10, 10], [3, 3, 10, 10]]:
             input_np = np.random.uniform(0.1, 0.5, dim).astype("float32")
@@ -298,6 +305,7 @@ def test_NNFunctionalMseLoss_sum(self):
             np.testing.assert_allclose(dy_result, expected, rtol=1e-05)
             self.assertEqual(dy_result.shape, ())
 
+    @test_with_pir_api
     def test_NNFunctionalMseLoss_none(self):
         for dim in [[10, 10], [2, 10, 10], [3, 3, 10, 10]]:
             input_np = np.random.uniform(0.1, 0.5, dim).astype("float32")
