fix pir ut test_adamw_op_(xpu) (#66982)

ooooo-create · web-flow · commit d798d293e7c6 · 2024-08-13T14:23:40.000+08:00
diff --git a/test/deprecated/legacy_test/CMakeLists.txt b/test/deprecated/legacy_test/CMakeLists.txt
@@ -690,7 +690,6 @@ set_tests_properties(test_conditional_block_deprecated
 
 # These UTs are to temporarily test static build for standalone_executor, will be removed after static build is enabled by default.
 set(STATIC_BUILD_TESTS
-    test_adamw_op_deprecated
     test_arg_min_max_op
     test_batch_norm_op_deprecated
     test_decoupled_py_reader_deprecated
diff --git a/test/deprecated/legacy_test/test_adamw_op_deprecated.py b/test/deprecated/legacy_test/test_adamw_op_deprecated.py
diff --git a/test/legacy_test/test_adamw_op.py b/test/legacy_test/test_adamw_op.py
@@ -231,6 +231,56 @@ def test_adamw_op_dygraph(self):
             adam.step()
             adam.clear_gradients()
 
+    def test_adamw_op(self):
+        paddle.enable_static()
+        place = base.CPUPlace()
+        shape = [2, 3, 8, 8]
+        exe = base.Executor(place)
+        train_prog = base.Program()
+        startup = base.Program()
+        with base.program_guard(train_prog, startup):
+            with base.unique_name.guard():
+                data = paddle.static.data(name="data", shape=shape)
+                conv = paddle.nn.Conv2D(
+                    in_channels=3,
+                    out_channels=8,
+                    kernel_size=3,
+                )(data)
+                loss = paddle.mean(conv)
+                if paddle.framework.in_pir_mode():
+                    beta1 = paddle.pir.core.create_persistable_value(
+                        shape=[1],
+                        dtype='float32',
+                        initializer=paddle.nn.initializer.Constant(0.85),
+                    )
+                    beta2 = paddle.pir.core.create_persistable_value(
+                        shape=[1],
+                        dtype='float32',
+                        initializer=paddle.nn.initializer.Constant(0.95),
+                    )
+                else:
+                    beta1 = paddle.static.create_global_var(
+                        shape=[1], value=0.85, dtype='float32', persistable=True
+                    )
+                    beta2 = paddle.static.create_global_var(
+                        shape=[1], value=0.95, dtype='float32', persistable=True
+                    )
+                betas = [beta1, beta2]
+                opt = paddle.optimizer.AdamW(
+                    learning_rate=1e-5,
+                    beta1=beta1,
+                    beta2=beta2,
+                    weight_decay=0.01,
+                    epsilon=1e-8,
+                )
+                opt.minimize(loss)
+
+        exe.run(startup)
+        data_np = np.random.random(shape).astype('float32')
+        rets = exe.run(train_prog, feed={"data": data_np}, fetch_list=[loss])
+        assert rets[0] is not None
+        paddle.disable_static()
+
     def test_adamw_op_dygraph_bypassing_step(self):
         paddle.disable_static()
         value = np.arange(26).reshape(2, 13).astype("float32")
diff --git a/test/xpu/test_adamw_op_xpu.py b/test/xpu/test_adamw_op_xpu.py
@@ -28,6 +28,7 @@
 import paddle
 import paddle.nn
 from paddle import base
+from paddle.pir_utils import test_with_pir_api
 
 
 def adamw_step(inputs, attributes):
@@ -201,6 +202,7 @@ def test_adamw_op_coverage(self):
             )
             assert adam.__str__() is not None
 
+        @test_with_pir_api
         def test_adamw_op(self):
             paddle.enable_static()
             place = base.XPUPlace(0)
@@ -218,18 +220,30 @@ def test_adamw_op(self):
                     )(data)
                     loss = paddle.mean(conv)
 
-                    beta1 = paddle.static.create_global_var(
-                        shape=[1],
-                        value=0.85,
-                        dtype=self.in_type_str,
-                        persistable=True,
-                    )
-                    beta2 = paddle.static.create_global_var(
-                        shape=[1],
-                        value=0.95,
-                        dtype=self.in_type_str,
-                        persistable=True,
-                    )
+                    if paddle.framework.in_pir_mode():
+                        beta1 = paddle.pir.core.create_persistable_value(
+                            shape=[1],
+                            dtype=self.in_type_str,
+                            initializer=paddle.nn.initializer.Constant(0.85),
+                        )
+                        beta2 = paddle.pir.core.create_persistable_value(
+                            shape=[1],
+                            dtype=self.in_type_str,
+                            initializer=paddle.nn.initializer.Constant(0.95),
+                        )
+                    else:
+                        beta1 = paddle.static.create_global_var(
+                            shape=[1],
+                            value=0.85,
+                            dtype=self.in_type_str,
+                            persistable=True,
+                        )
+                        beta2 = paddle.static.create_global_var(
+                            shape=[1],
+                            value=0.95,
+                            dtype=self.in_type_str,
+                            persistable=True,
+                        )
                     betas = [beta1, beta2]
                     opt = paddle.optimizer.AdamW(
                         learning_rate=1e-5,
@@ -457,6 +471,7 @@ def get_numpy_output(param, grad, moment1, moment2, lr_ratio, t):
                     linear2.bias.numpy(), fc2_b, rtol=1e-5, atol=1e-5
                 )
 
+        @test_with_pir_api
         def test_adamw_op(self):
             paddle.enable_static()
             place = base.XPUPlace(0)
@@ -536,7 +551,7 @@ def test_adamw_op(self):
                         epsilon=epsilon,
                         lr_ratio=simple_lr_fun,
                     )
-                    opt.minimize(avg_cost)
+                    _, params_grads = opt.minimize(avg_cost)
 
             def get_numpy_output(param, grad, moment1, moment2, lr_ratio, t):
                 np_inputs = {
@@ -562,26 +577,44 @@ def get_numpy_output(param, grad, moment1, moment2, lr_ratio, t):
                 )
                 return param_out, moment1_out, moment2_out
 
-            fetch_list1 = [
-                "linear_0.w_0",
-                "linear_0.b_0",
-                "linear_1.w_0",
-                "linear_1.b_0",
-            ]
-            fetch_list2 = [
-                "linear_0.w_0",
-                "linear_0.w_0@GRAD",
-                "linear_0.b_0",
-                "linear_0.b_0@GRAD",
-                "linear_1.w_0",
-                "linear_1.w_0@GRAD",
-                "linear_1.b_0",
-                "linear_1.b_0@GRAD",
-            ]
-
+            if paddle.framework.in_pir_mode():
+                fetch_list1 = [
+                    linear1.weight,
+                    linear1.bias,
+                    linear2.weight,
+                    linear2.bias,
+                ]
+                fetch_list2 = []
+                for param in fetch_list1:
+                    for item, vale in dict(params_grads).items():
+                        if param.is_same(item):
+                            fetch_list2.append(item)
+                            fetch_list2.append(vale)
+            else:
+                fetch_list1 = [
+                    "linear_0.w_0",
+                    "linear_0.b_0",
+                    "linear_1.w_0",
+                    "linear_1.b_0",
+                ]
+                fetch_list2 = [
+                    "linear_0.w_0",
+                    "linear_0.w_0@GRAD",
+                    "linear_0.b_0",
+                    "linear_0.b_0@GRAD",
+                    "linear_1.w_0",
+                    "linear_1.w_0@GRAD",
+                    "linear_1.b_0",
+                    "linear_1.b_0@GRAD",
+                ]
             exe = base.Executor(place)
             exe.run(startup)
-            test_prog = train_prog.clone(for_test=True)
+            if paddle.framework.in_pir_mode():
+                value_map = paddle.pir.IrMapping()
+                test_prog = train_prog.clone(value_map)
+                fetch_list1 = [value_map.look_up(v) for v in fetch_list1]
+            else:
+                test_prog = train_prog.clone(for_test=True)
 
             for i in range(5):
                 inputs = np.random.random(size=[8, 10]).astype('float32')
