diff --git a/python/paddle/incubate/optimizer/lookahead.py b/python/paddle/incubate/optimizer/lookahead.py index 821b5c3ce036c..12cb00ba7a3ff 100644 --- a/python/paddle/incubate/optimizer/lookahead.py +++ b/python/paddle/incubate/optimizer/lookahead.py @@ -17,7 +17,9 @@ from paddle.base.dygraph import base as imperative_base from paddle.base.framework import Variable from paddle.base.layer_helper import LayerHelper +from paddle.framework import in_pir_mode from paddle.optimizer import Optimizer +from paddle.pir.core import create_parameter __all__ = [] @@ -121,7 +123,9 @@ def __init__(self, inner_optimizer, alpha=0.5, k=5, name=None): self.inner_optimizer = inner_optimizer if self.inner_optimizer._parameter_list is None: parameters = ( - framework.default_main_program().global_block().all_parameters() + paddle.static.default_main_program() + .global_block() + .all_parameters() ) else: parameters = self.inner_optimizer._parameter_list @@ -186,40 +190,64 @@ def step(self): ) def _create_accumulators(self, block, parameters): - assert isinstance(block, framework.Block) + assert isinstance(block, (framework.Block, paddle.pir.Block)) for p in parameters: self._add_accumulator(self._slow_str, p) def _increment_global_var(self): - if self._global_step_var is None: - self._global_step_var = paddle.static.create_global_var( - name=unique_name.generate("lookahead_step"), - shape=[1], - value=0, - dtype='int32', - persistable=True, + if in_pir_mode(): + if self._global_step_var is None: + self._global_step_var = create_parameter( + dtype='int32', + shape=[1], + name=unique_name.generate("lookahead_step"), + trainable=False, + initializer=paddle.nn.initializer.ConstantInitializer( + value=0.0, force_cpu=False + ), + ) + self._global_step_var = paddle.increment(self._global_step_var, 1.0) + else: + if self._global_step_var is None: + self._global_step_var = paddle.static.create_global_var( + name=unique_name.generate("lookahead_step"), + shape=[1], + value=0, + dtype='int32', + persistable=True, + ) + + self.helper.append_op( + type='increment', + inputs={'X': [self._global_step_var]}, + outputs={'Out': [self._global_step_var]}, + attrs={'step': 1.0}, ) - self.helper.append_op( - type='increment', - inputs={'X': [self._global_step_var]}, - outputs={'Out': [self._global_step_var]}, - attrs={'step': 1.0}, - ) - def _append_optimize_op(self, block, param_and_grad): one_var = paddle.ones(shape=[1], dtype='int32', name='lookahead_ones') zero_var = paddle.zeros( shape=[1], dtype='int32', name='lookahead_zeros' ) - k_var = paddle.static.create_global_var( - name=unique_name.generate("lookahead_k"), - shape=[1], - value=self.k, - dtype='int32', - persistable=True, - ) + if in_pir_mode(): + k_var = create_parameter( + dtype='int32', + shape=[1], + name=unique_name.generate("lookahead_k"), + trainable=False, + initializer=paddle.nn.initializer.ConstantInitializer( + value=float(self.k), force_cpu=False + ), + ) + else: + k_var = paddle.static.create_global_var( + name=unique_name.generate("lookahead_k"), + shape=[1], + value=self.k, + dtype='int32', + persistable=True, + ) mod = paddle.remainder(self._global_step_var, k_var) @@ -284,7 +312,9 @@ def minimize( >>> lookahead.clear_grad() """ - assert isinstance(loss, Variable), "The loss should be an Tensor." + assert isinstance( + loss, (Variable, paddle.pir.Value) + ), "The loss should be an Tensor." # Apply inner optimizer to the main_program optimize_ops, params_grads = self.inner_optimizer.minimize( diff --git a/python/paddle/optimizer/sgd.py b/python/paddle/optimizer/sgd.py index dcd8e17d51cb6..233261f7f769b 100644 --- a/python/paddle/optimizer/sgd.py +++ b/python/paddle/optimizer/sgd.py @@ -14,7 +14,7 @@ import warnings -from paddle import _C_ops +from paddle import _C_ops, pir from ..base import framework from ..base.dygraph import no_grad @@ -93,7 +93,7 @@ def __init__( self._master_weights = {} def _create_accumulators(self, block, parameters): - assert isinstance(block, framework.Block) + assert isinstance(block, (framework.Block, pir.Block)) if isinstance(parameters, dict): parameters = self._update_param_group(parameters) diff --git a/test/legacy_test/test_lookahead.py b/test/legacy_test/test_lookahead.py index f9eeee2f651a4..911db37cd8977 100644 --- a/test/legacy_test/test_lookahead.py +++ b/test/legacy_test/test_lookahead.py @@ -18,6 +18,8 @@ import paddle from paddle import base, nn +from paddle.base.framework import in_pir_mode +from paddle.pir_utils import test_with_pir_api LOOKAHEAD_K = 5 LOOKAHEAD_ALPHA = 0.2 @@ -25,20 +27,21 @@ class TestLookAhead(unittest.TestCase): + @test_with_pir_api def test_lookahead_static(self): paddle.enable_static() place = base.CPUPlace() shape = [2, 3, 8, 8] exe = base.Executor(place) - train_program = base.Program() - startup = base.Program() - with base.program_guard(train_program, startup): + train_program = paddle.static.Program() + startup = paddle.static.Program() + with paddle.static.program_guard(train_program, startup): with base.unique_name.guard(): data = paddle.static.data( name='X', shape=[None, 1], dtype='float32' ) - hidden = paddle.static.nn.fc(x=data, size=10) - loss = paddle.mean(hidden) + hidden = paddle.nn.Linear(1, 10) + loss = paddle.mean(hidden(data)) optimizer = paddle.optimizer.SGD(learning_rate=SGD_LR) lookahead = paddle.incubate.optimizer.LookAhead( @@ -55,13 +58,20 @@ def test_lookahead_static(self): fast_param - slow_param ) x = np.random.random(size=(10, 1)).astype('float32') + if in_pir_mode(): + for op in train_program.global_block().ops: + if op.name() == 'pd_op.add_grad': + bias_grad = op.result(1) + fetch_list = [hidden.bias, bias_grad] + else: + fetch_list = [ + 'linear_0.b_0', + 'linear_0.b_0@GRAD', + ] latest_b, b_grad = exe.run( program=train_program, feed={'X': x}, - fetch_list=[ - 'fc_0.b_0', - 'fc_0.b_0@GRAD', - ], + fetch_list=fetch_list, ) if i == 0: slow_param = latest_b