From ea60e64470fdec696128a1f7cc5c83f59c233d43 Mon Sep 17 00:00:00 2001
From: mapingshuo <mps2012@yeah.net>
Date: Fri, 24 Jul 2020 14:04:11 +0800
Subject: [PATCH] correct the LookaheadOptimizer programDesc, test=develop
 (#25688)

---
 python/paddle/fluid/optimizer.py | 88 +++++++++++++++++---------------
 1 file changed, 46 insertions(+), 42 deletions(-)

diff --git a/python/paddle/fluid/optimizer.py b/python/paddle/fluid/optimizer.py
index 165c44b964..a8a7f217bd 100644
--- a/python/paddle/fluid/optimizer.py
+++ b/python/paddle/fluid/optimizer.py
@@ -4884,48 +4884,52 @@ class LookaheadOptimizer(object):
                 inputs={"X": fast_var},
                 outputs={"Out": slow_var})
 
-        # Add Var k to main prog and startup prog
-        k = layers.create_global_var(
-            name="lookahead_k",
-            shape=[1],
-            value=int(self.k),
-            dtype='int32',
-            persistable=True)
+        with framework.program_guard(main_block.program, startup_program):
+            # Add Var k to main prog and startup prog
+            k = layers.create_global_var(
+                name="lookahead_k",
+                shape=[1],
+                value=int(self.k),
+                dtype='int32',
+                persistable=True)
 
-        # Add Var alpha to main prog and startup prog
-        alpha = layers.create_global_var(
-            name="lookahead_alpha",
-            shape=[1],
-            value=float(self.alpha),
-            dtype='float32',
-            persistable=True)
+            # Add Var alpha to main prog and startup prog
+            alpha = layers.create_global_var(
+                name="lookahead_alpha",
+                shape=[1],
+                value=float(self.alpha),
+                dtype='float32',
+                persistable=True)
 
-        # Add Var step
-        step = layers.create_global_var(
-            name="lookahead_step",
-            shape=[1],
-            value=int(0),
-            dtype='int32',
-            persistable=True)
-        layers.increment(x=step, value=1.0, in_place=True)
-
-        # lookahead
-        zero_var = layers.fill_constant(shape=[1], dtype='float32', value=0.0)
-
-        one_var = layers.fill_constant(shape=[1], dtype='float32', value=1.0)
-
-        mod = layers.elementwise_mod(step, k)
-        with layers.control_flow.Switch() as switch:
-            with switch.case(mod == zero_var):
-                for param_name in params:
-                    fast_var = main_block.var(param_name)
-                    slow_var = param_to_slow[param_name]
-                    tmp_var = layers.elementwise_add(
-                        layers.elementwise_mul(fast_var, alpha),
-                        layers.elementwise_mul(
-                            slow_var, layers.elementwise_sub(one_var, alpha)))
-                    layers.assign(input=tmp_var, output=slow_var)
-                    layers.assign(input=tmp_var, output=fast_var)
-            with switch.default():
-                pass
+            # Add Var step
+            step = layers.create_global_var(
+                name="lookahead_step",
+                shape=[1],
+                value=int(0),
+                dtype='int32',
+                persistable=True)
+            layers.increment(x=step, value=1.0, in_place=True)
+
+            # lookahead
+            zero_var = layers.fill_constant(
+                shape=[1], dtype='float32', value=0.0)
+
+            one_var = layers.fill_constant(
+                shape=[1], dtype='float32', value=1.0)
+
+            mod = layers.elementwise_mod(step, k)
+            with layers.control_flow.Switch() as switch:
+                with switch.case(mod == zero_var):
+                    for param_name in params:
+                        fast_var = main_block.var(param_name)
+                        slow_var = param_to_slow[param_name]
+                        tmp_var = layers.elementwise_add(
+                            layers.elementwise_mul(fast_var, alpha),
+                            layers.elementwise_mul(
+                                slow_var,
+                                layers.elementwise_sub(one_var, alpha)))
+                        layers.assign(input=tmp_var, output=slow_var)
+                        layers.assign(input=tmp_var, output=fast_var)
+                with switch.default():
+                    pass
         return mini_out
-- 
GitLab