【PIR API adaptor No.272】Migrate LookAhead to pir

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(

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)
 

test/legacy_test/test_lookahead.py

@@ -18,27 +18,30 @@
 
 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
 SGD_LR = 1.0
 
 
 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