[PIR]Migrate Momentum into pir (#58730)

0x45f · web-flow · commit 54ac5b556be9 · 2023-11-09T16:23:09.000+08:00
diff --git a/python/paddle/nn/functional/loss.py b/python/paddle/nn/functional/loss.py
@@ -22,7 +22,11 @@
 from paddle.utils import deprecated
 
 from ...base.data_feeder import check_variable_and_dtype
-from ...base.framework import _current_expected_place, in_pir_mode
+from ...base.framework import (
+    _current_expected_place,
+    in_dynamic_or_pir_mode,
+    in_pir_mode,
+)
 from ...base.layer_helper import LayerHelper
 from ...common_ops_import import Variable
 from ...tensor.manipulation import reshape
@@ -420,7 +424,7 @@ def square_error_cost(input, label):
                     [0.01000000, 0.01000000])
 
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         minus_out = _C_ops.subtract(input, label)
         square_out = _C_ops.square(minus_out)
         return square_out
diff --git a/python/paddle/optimizer/momentum.py b/python/paddle/optimizer/momentum.py
@@ -16,7 +16,7 @@
 
 import paddle
 from paddle import _C_ops
-from paddle.framework import in_dynamic_mode
+from paddle.framework import in_dynamic_or_pir_mode
 from paddle.regularizer import L2Decay
 
 from ..base import core, framework
@@ -204,7 +204,7 @@ def _create_accumulators(self, block, parameters):
         if framework.in_dynamic_mode():
             return
         '''
-        assert isinstance(block, framework.Block)
+        assert isinstance(block, (framework.Block, paddle.pir.Block))
 
         if isinstance(parameters, dict):
             parameters = self._update_param_group(parameters)
@@ -276,7 +276,7 @@ def _append_optimize_op(self, block, param_and_grad):
             else None
         )
 
-        if in_dynamic_mode():
+        if in_dynamic_or_pir_mode():
             if isinstance(param_and_grad, dict):
                 self._update_regularization(param_and_grad['weight_decay'])
             return _C_ops.momentum_(
@@ -472,13 +472,17 @@ def _append_optimize_multi_tensor_op(
                     else None
                 )
 
-                if in_dynamic_mode():
+                if in_dynamic_or_pir_mode():
                     found_inf = self._get_auxiliary_var('found_inf')
                     if found_inf:
-                        if isinstance(found_inf, core.eager.Tensor):
+                        if isinstance(
+                            found_inf, (core.eager.Tensor, paddle.pir.OpResult)
+                        ):
                             self._set_auxiliary_var('found_inf', True)
                     else:
-                        if isinstance(found_inf, core.eager.Tensor):
+                        if isinstance(
+                            found_inf, (core.eager.Tensor, paddle.pir.OpResult)
+                        ):
                             self._set_auxiliary_var('found_inf', False)
                         _, _, _ = _C_ops.merged_momentum_(
                             self._param_dict[key][param_group_idx],
diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py
@@ -846,50 +846,61 @@ def _add_accumulator(
             )
         if shape is None:
             shape = param.shape
-        assert isinstance(self.helper, LayerHelper)
 
         var_name = param.name + "_" + name
         var_name = unique_name.generate(var_name)
         self._opti_name_list.append(var_name)
 
-        var = self.helper.create_global_variable(
-            name=var_name,
-            persistable=True,
-            dtype=dtype or param.dtype,
-            type=core.VarDesc.VarType.LOD_TENSOR,
-            shape=shape,
-            belong_to_optimizer=True,
-        )
         if device is None:
             device = self._get_device_for_param(param.name)
 
-        if (
-            in_dygraph_mode()
-            and (device == 'cpu' or isinstance(device, core.CPUPlace))
-            and (not core.is_compiled_with_xpu())
-        ):
-            _C_ops.full_(
-                var,
-                var.shape,
-                str(float(fill_value)),
-                var.dtype,
-                core.CPUPlace(),
+        if in_pir_mode():
+            var = paddle.pir.core.create_parameter(
+                dtype or param.dtype,
+                shape,
+                var_name,
+                initializer=paddle.nn.initializer.Constant(
+                    value=float(fill_value)
+                ),
             )
         else:
-            with device_guard(device):
-                self.helper.set_variable_initializer(
+            assert isinstance(self.helper, LayerHelper)
+            var = self.helper.create_global_variable(
+                name=var_name,
+                persistable=True,
+                dtype=dtype or param.dtype,
+                type=core.VarDesc.VarType.LOD_TENSOR,
+                shape=shape,
+                belong_to_optimizer=True,
+            )
+
+            if (
+                in_dygraph_mode()
+                and (device == 'cpu' or isinstance(device, core.CPUPlace))
+                and (not core.is_compiled_with_xpu())
+            ):
+                _C_ops.full_(
                     var,
-                    initializer=paddle.nn.initializer.Constant(
-                        value=float(fill_value)
-                    ),
+                    var.shape,
+                    str(float(fill_value)),
+                    var.dtype,
+                    core.CPUPlace(),
                 )
+            else:
+                with device_guard(device):
+                    self.helper.set_variable_initializer(
+                        var,
+                        initializer=paddle.nn.initializer.Constant(
+                            value=float(fill_value)
+                        ),
+                    )
 
-        if framework.in_dygraph_mode():
-            if len(self._accumulators_holder) > 0:
-                assert (
-                    var_name in self._accumulators_holder
-                ), f"Optimizer set error, {var_name} should in state dict"
-                var.set_value(self._accumulators_holder.pop(var_name))
+            if framework.in_dygraph_mode():
+                if len(self._accumulators_holder) > 0:
+                    assert (
+                        var_name in self._accumulators_holder
+                    ), f"Optimizer set error, {var_name} should in state dict"
+                    var.set_value(self._accumulators_holder.pop(var_name))
 
         self._accumulators[name][param.name] = var
         return var
diff --git a/python/paddle/pir/core.py b/python/paddle/pir/core.py
@@ -260,6 +260,7 @@ def __init__(self, shape, dtype):
 def create_parameter(
     dtype,
     shape,
+    name=None,
     **kwargs,
 ):
     if 'initializer' not in kwargs:
@@ -269,7 +270,9 @@ def create_parameter(
     if dtype is not None:
         if not isinstance(dtype, DataType):
             dtype = convert_np_dtype_to_dtype_(dtype)
-    op_result_name = unique_name.generate('parameter')
+    op_result_name = name
+    if not op_result_name:
+        op_result_name = unique_name.generate('parameter')
     startup_program = default_startup_program()
     main_program = default_main_program()
     parameter_meta = ParameterMeta(shape, dtype)
diff --git a/python/paddle/regularizer.py b/python/paddle/regularizer.py
@@ -13,9 +13,9 @@
 # limitations under the License.
 
 
-from paddle import _C_ops
+from paddle import _C_ops, pir
 from paddle.base import framework
-from paddle.base.framework import in_dygraph_mode
+from paddle.base.framework import in_dynamic_or_pir_mode
 
 __all__ = ['L1Decay', 'L2Decay']
 
@@ -123,10 +123,12 @@ def __call__(self, param, grad, block):
         Returns:
             new variable for weight decay
         """
-        assert isinstance(param, framework.Variable)
-        assert isinstance(block, framework.Block)
+        assert isinstance(
+            param, (framework.Variable, pir.OpResult, pir.core.ParameterMeta)
+        )
+        assert isinstance(block, (framework.Block, pir.Block))
 
-        if in_dygraph_mode():
+        if in_dynamic_or_pir_mode():
             sign = _C_ops.sign(param)
             return _C_ops.scale(sign, self._coeff, 0.0, True)
         else:
@@ -233,10 +235,12 @@ def __call__(self, param, grad, block):
         Returns:
             new variable for weight decay
         """
-        assert isinstance(param, framework.Variable)
-        assert isinstance(block, framework.Block)
+        assert isinstance(
+            param, (framework.Variable, pir.OpResult, pir.core.ParameterMeta)
+        )
+        assert isinstance(block, (framework.Block, pir.Block))
 
-        if in_dygraph_mode():
+        if in_dynamic_or_pir_mode():
             return _C_ops.scale(param, self._coeff, 0.0, True)
         else:
             decay = block.create_var(
diff --git a/python/paddle/static/input.py b/python/paddle/static/input.py
@@ -23,6 +23,7 @@
     static_only,
 )
 from paddle.base.layer_helper import LayerHelper
+from paddle.base.libpaddle import DataType
 
 from ..base.variable_index import _setitem_impl_, _setitem_static
 
@@ -126,7 +127,9 @@ def _reset_data_op_insertion_point():
         dtype = paddle.get_default_dtype()
 
     if in_pir_mode():
-        ir_dtype = paddle.pir.core.convert_np_dtype_to_dtype_(dtype)
+        ir_dtype = dtype
+        if not isinstance(ir_dtype, DataType):
+            ir_dtype = paddle.pir.core.convert_np_dtype_to_dtype_(dtype)
         _reset_data_op_insertion_point()
         out = paddle._pir_ops.data(name, shape, ir_dtype, core.Place())
         paddle.pir.reset_insertion_point_to_end()
diff --git a/test/legacy_test/test_momentum_op.py b/test/legacy_test/test_momentum_op.py
@@ -22,6 +22,7 @@
 import paddle
 from paddle import base
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 
 def calculate_momentum_by_numpy(
@@ -120,15 +121,15 @@ def init_dtype(self):
         pass
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
 
 class TestMomentumOpFp16(TestMomentumOp1):
     def init_dtype(self):
         self.dtype = np.float16
 
     def test_check_output(self):
-        self.check_output(atol=1e-3)
+        self.check_output(atol=1e-3, check_pir=True)
 
 
 class TestMomentumOp2(OpTest):
@@ -166,7 +167,7 @@ def setUp(self):
         self.outputs = {'ParamOut': param_out, 'VelocityOut': velocity_out}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
 
 @unittest.skipIf(
@@ -553,14 +554,17 @@ def test_momentum_dygraph(self):
         adam.step()
         adam.clear_gradients()
 
+    @test_with_pir_api
     def test_momentum(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, activation=None)
+            linear = paddle.nn.Linear(13, 1)
+            y_predict = linear(x)
             cost = paddle.nn.functional.square_error_cost(
                 input=y_predict, label=y
             )
@@ -575,11 +579,17 @@ def test_momentum(self):
             train_reader = paddle.batch(
                 paddle.dataset.uci_housing.train(), batch_size=1
             )
-            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)
+                exe.run(
+                    main,
+                    feed={
+                        'x': data[0][0].astype('float32'),
+                        'y': data[0][1].astype('float32'),
+                    },
+                    fetch_list=fetch_list,
+                )
 
     def test_raise_error(self):
         self.assertRaises(
@@ -639,7 +649,7 @@ def init_config(self):
 
     def test_check_output(self):
         paddle.enable_static()
-        self.check_output()
+        self.check_output(check_pir=True)
 
 
 class TestMomentumOpWithDecayFP16(TestMomentumOpWithDecay):
@@ -648,7 +658,7 @@ def init_config(self):
 
     def test_check_output(self):
         paddle.enable_static()
-        self.check_output(atol=1e-3)
+        self.check_output(atol=1e-3, check_pir=True)
 
 
 class TestMomentumOpWithDecay2(TestMomentumOpWithDecay):
@@ -690,14 +700,17 @@ def test_momentum_dygraph_1(self):
             regularization=paddle.regularizer.L2Decay(coeff=0.1)
         )
 
+    @test_with_pir_api
     def test_momentum_static(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, activation=None)
+            linear = paddle.nn.Linear(13, 1)
+            y_predict = linear(x)
             cost = paddle.nn.functional.square_error_cost(
                 input=y_predict, label=y
             )
@@ -712,11 +725,17 @@ def test_momentum_static(self):
             train_reader = paddle.batch(
                 paddle.dataset.uci_housing.train(), batch_size=1
             )
-            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)
+                exe.run(
+                    main,
+                    feed={
+                        'x': data[0][0].astype('float32'),
+                        'y': data[0][1].astype('float32'),
+                    },
+                    fetch_list=fetch_list,
+                )
 
 
 class TestFusedMomentumWithDecayAPI(unittest.TestCase):
