【PIR API adaptor No.124】 optimizer.Lamb (#58881)

Liyulingyue · web-flow · commit 39693325e01d · 2023-12-25T14:28:02.000+08:00
diff --git a/python/paddle/optimizer/lamb.py b/python/paddle/optimizer/lamb.py
@@ -243,7 +243,7 @@ def _append_optimize_op(self, block, param_and_grad):
         else:
             master_weight = None
 
-        if framework.in_dygraph_mode():
+        if framework.in_dynamic_or_pir_mode():
             _C_ops.lamb_(
                 param_and_grad[0],
                 param_and_grad[1],
diff --git a/python/paddle/tensor/linalg.py b/python/paddle/tensor/linalg.py
@@ -428,7 +428,7 @@ def vector_norm(
           name (str, optional): The default value is None. Normally there is no need for
               user to set this property. For more information, please refer to :ref:`api_guide_Name`.
         """
-        if in_dynamic_mode():
+        if in_dynamic_or_pir_mode():
             if axis is None:
                 axis = -1
             return _C_ops.p_norm(input, porder, axis, 1e-12, keepdim, asvector)
diff --git a/test/legacy_test/test_imperative_base.py b/test/legacy_test/test_imperative_base.py
@@ -14,15 +14,15 @@
 
 import contextlib
 
-from paddle import base
+from paddle import base, static
 
 
 @contextlib.contextmanager
 def new_program_scope(main=None, startup=None, scope=None):
-    prog = main if main else base.Program()
-    startup_prog = startup if startup else base.Program()
+    prog = main if main else static.Program()
+    startup_prog = startup if startup else static.Program()
     scope = scope if scope else base.core.Scope()
-    with base.scope_guard(scope):
-        with base.program_guard(prog, startup_prog):
+    with static.scope_guard(scope):
+        with static.program_guard(prog, startup_prog):
             with base.unique_name.guard():
                 yield
diff --git a/test/legacy_test/test_imperative_optimizer_v2.py b/test/legacy_test/test_imperative_optimizer_v2.py
@@ -22,6 +22,7 @@
 from paddle import base
 from paddle.base import core
 from paddle.distributed.fleet.meta_optimizers import DGCMomentumOptimizer
+from paddle.pir_utils import test_with_pir_api
 
 # Note(wangzhongpu)
 # In dygraph, don't support ModelAverage, DGCMomentumOptimizer, ExponentialMovingAverage, PipelineOptimizer, LookaheadOptimizer, RecomputeOptimizer.
@@ -790,6 +791,7 @@ def get_optimizer(self):
         return optimizer
 
     # should fix: may fail in CI-windows
+    @test_with_pir_api
     def _test_lamb(self):
         self._check_mlp()
 
diff --git a/test/legacy_test/test_lambv2_op.py b/test/legacy_test/test_lambv2_op.py
@@ -20,6 +20,7 @@
 from paddle import base
 from paddle.base import core
 from paddle.base.dygraph.base import switch_to_static_graph
+from paddle.pir_utils import test_with_pir_api
 
 
 class LAMBOptimizer(paddle.optimizer.Lamb):
@@ -113,6 +114,7 @@ def test_lamb_op(self):
 
 
 class TestLambOpWithCombinedOp(unittest.TestCase):
+    @test_with_pir_api
     def test_lamb_op_with_multi_steps(self):
         paddle.enable_static()
 
@@ -124,7 +126,10 @@ def _build_static_model(main, startup, seed=100):
                     name='X', shape=[-1, 13], dtype='float32'
                 )
                 y = paddle.static.data(name='Y', shape=[-1, 1], dtype='float32')
-                prediction = paddle.static.nn.fc(x, size=1, activation=None)
+                linear = paddle.nn.Linear(
+                    in_features=x.shape[-1], out_features=1
+                )
+                prediction = linear(x)
                 loss = paddle.nn.functional.square_error_cost(
                     input=prediction, label=y
                 )
@@ -138,8 +143,8 @@ def _build_static_model(main, startup, seed=100):
             feed_x = np.random.random(size=(10, 13)).astype('float32')
             feed_y = np.random.random(size=(10, 1)).astype('float32')
 
-            main_program = base.Program()
-            startup_program = base.Program()
+            main_program = paddle.static.Program()
+            startup_program = paddle.static.Program()
             with base.program_guard(main_program, startup_program):
                 avg_loss = _build_static_model(main_program, startup_program)
                 lamb_kernel = paddle.optimizer.Lamb(learning_rate=0.2)
@@ -150,11 +155,11 @@ def _build_static_model(main, startup, seed=100):
             output = executor.run(
                 program=main_program,
                 feed={'X': feed_x, 'Y': feed_y},
-                fetch_list=[avg_loss.name],
+                fetch_list=[avg_loss],
             )
 
-            main = base.Program()
-            startup = base.Program()
+            main = paddle.static.Program()
+            startup = paddle.static.Program()
             with base.program_guard(main, startup):
                 loss = _build_static_model(main, startup)
                 lamb = LAMBOptimizer(learning_rate=0.2)
@@ -165,7 +170,7 @@ def _build_static_model(main, startup, seed=100):
             out = exe.run(
                 program=main,
                 feed={'X': feed_x, 'Y': feed_y},
-                fetch_list=[loss.name],
+                fetch_list=[loss],
             )
 
             np.testing.assert_allclose(out, output, rtol=1e-05)
@@ -226,8 +231,7 @@ def check_main(self, x_np, place, multi_precision=False, seed=10, n=10):
         weight, bias = linear.weight, linear.bias
         exe = paddle.static.Executor(place)
         scope = paddle.static.Scope()
-        x = main_prog.global_block().var(x.name)
-        if x.dtype == core.VarDesc.VarType.FP16:
+        if x.dtype in (core.VarDesc.VarType.FP16, core.DataType.FLOAT16):
             x_np = x_np.astype(np.float16)
 
         def get_parameter(var):
@@ -256,7 +260,7 @@ def get_parameter(var):
 
             weight_np, bias_np = None, None
             for i in range(n):
-                feed_dict = {x.name: x_np}
+                feed_dict = {'x': x_np}
                 weight_np, bias_np = exe.run(
                     main_prog, feed=feed_dict, fetch_list=[weight, bias]
                 )
