【PIR API adaptor No.182】Migrate paddle.optimizer.RMSProp into pir (#60350)

ooooo-create · web-flow · commit 877a1c881297 · 2024-01-23T19:24:54.000+08:00
diff --git a/python/paddle/optimizer/rmsprop.py b/python/paddle/optimizer/rmsprop.py
@@ -17,7 +17,7 @@
 from paddle import _C_ops
 
 from ..base import framework
-from ..base.framework import in_dygraph_mode
+from ..base.framework import in_dynamic_or_pir_mode
 from .optimizer import Optimizer
 
 __all__ = []
@@ -252,7 +252,7 @@ def _append_optimize_op(self, block, param_and_grad):
             else None
         )
 
-        if in_dygraph_mode():
+        if in_dynamic_or_pir_mode():
             _C_ops.rmsprop_(
                 param_and_grad[0],
                 mean_square_acc,
diff --git a/test/legacy_test/test_rmsprop_op.py b/test/legacy_test/test_rmsprop_op.py
@@ -20,6 +20,7 @@
 import paddle
 from paddle import base
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 
 def create_selected_rows_and_tensor(
@@ -272,14 +273,18 @@ def test_rmsprop_dygraph(self):
         adam.step()
         adam.clear_gradients()
 
+    @test_with_pir_api
     def test_rmsprop(self):
         paddle.enable_static()
         place = base.CPUPlace()
-        main = base.Program()
-        with base.program_guard(main):
+        main = paddle.static.Program()
+        startup = paddle.static.Program()
+        with paddle.static.program_guard(main, startup):
             x = paddle.static.data(name='x', shape=[-1, 13], dtype='float32')
             y = paddle.static.data(name='y', shape=[-1, 1], dtype='float32')
-            y_predict = paddle.static.nn.fc(x, size=1)
+            y_predict = paddle.nn.Linear(
+                in_features=x.shape[-1], out_features=1
+            )(x)
             cost = paddle.nn.functional.square_error_cost(
                 input=y_predict, label=y
             )
@@ -294,7 +299,7 @@ def test_rmsprop(self):
             )
             feeder = base.DataFeeder(place=place, feed_list=[x, y])
             exe = base.Executor(place)
-            exe.run(base.default_startup_program())
+            exe.run(startup)
             for data in train_reader():
                 exe.run(main, feed=feeder.feed(data), fetch_list=fetch_list)
 
@@ -468,7 +473,9 @@ def static_rmsprop_mp(self, mp, use_amp):
                 data = paddle.static.data(
                     shape=[2, 2], name='X', dtype='float32'
                 )
-            hidden = paddle.static.nn.fc(x=data, size=10)
+            hidden = paddle.nn.Linear(
+                in_features=data.shape[-1], out_features=10
+            )(data)
             loss = paddle.mean(hidden)
             optimizer.minimize(loss)
         exe.run(startup_program)
@@ -483,11 +490,61 @@ def static_rmsprop_mp(self, mp, use_amp):
         out = []
         for idx in range(5):
             (loss_data,) = exe.run(
-                train_program, feed={"X": x}, fetch_list=[loss.name]
+                train_program, feed={"X": x}, fetch_list=[loss]
             )
             out.append(loss_data)
         return out
 
+    def pir_rmsprop_mp(self, mp, use_amp):
+        with paddle.pir_utils.IrGuard():
+            paddle.seed(100)
+            np.random.seed(100)
+            exe = paddle.static.Executor('gpu')
+            train_program = paddle.static.Program()
+            startup_program = paddle.static.Program()
+            optimizer = paddle.optimizer.RMSProp(0.1)
+            optimizer._multi_precision = mp
+
+            if use_amp:
+                optimizer = paddle.static.amp.decorate(
+                    optimizer,
+                    init_loss_scaling=128.0,
+                    use_dynamic_loss_scaling=True,
+                    use_pure_fp16=True,
+                    use_fp16_guard=False,
+                )
+            with paddle.static.program_guard(train_program, startup_program):
+                if use_amp:
+                    data = paddle.static.data(
+                        shape=[2, 2], name='X', dtype='float16'
+                    )
+                else:
+                    data = paddle.static.data(
+                        shape=[2, 2], name='X', dtype='float32'
+                    )
+                hidden = paddle.nn.Linear(
+                    in_features=data.shape[-1], out_features=10
+                )(data)
+                loss = paddle.mean(hidden)
+                optimizer.minimize(loss)
+            exe.run(startup_program)
+
+            if use_amp:
+                optimizer.amp_init(
+                    place=paddle.CUDAPlace(0),
+                    scope=paddle.static.global_scope(),
+                )
+                x = np.random.random(size=(2, 2)).astype('float16')
+            else:
+                x = np.random.random(size=(2, 2)).astype('float32')
+            out = []
+            for idx in range(5):
+                (loss_data,) = exe.run(
+                    train_program, feed={"X": x}, fetch_list=[loss]
+                )
+                out.append(loss_data)
+            return out
+
     def test_main(self):
         if not paddle.is_compiled_with_cuda():
             return
@@ -519,6 +576,16 @@ def test_main(self):
                 rtol=1e-05,
                 atol=0.1,
             )
+        # NOT support amp training "Test pir mode"
+        # output1_pir = self.pir_rmsprop_mp(use_amp=True, mp=True)
+        # output2_pir = self.pir_rmsprop_mp(use_amp=False, mp=False)
+        # for idx in range(len(output1_pir)):
+        #     np.testing.assert_allclose(
+        #         output1_pir[idx].astype('float32'),
+        #         output2_pir[idx].astype('float32'),
+        #         rtol=1e-05,
+        #         atol=0.1,
+        #     )
 
 
 if __name__ == "__main__":
