Merge pull request #12254 from JiayiFeng/fix_lr_decay

Fix learning rate scheduler performance issue

python/paddle/fluid/layers/learning_rate_scheduler.py

@@ -62,7 +62,7 @@ def noam_decay(d_model, warmup_steps):
         The decayed learning rate.
     """
     global_step = _decay_step_counter(1)
-    with init_on_cpu():
+
     a = global_step**-0.5
     b = (warmup_steps**-1.5) * global_step
     lr_value = (d_model**-0.5) * ops.elementwise_min(a, b)
@@ -108,8 +108,6 @@ def exponential_decay(learning_rate, decay_steps, decay_rate, staircase=False):
     """
     global_step = _decay_step_counter()
 
-    with init_on_cpu():
-        # update learning_rate
     div_res = global_step / decay_steps
     if staircase:
         div_res = ops.floor(div_res)
@@ -138,7 +136,6 @@ def natural_exp_decay(learning_rate, decay_steps, decay_rate, staircase=False):
     """
     global_step = _decay_step_counter()
 
-    with init_on_cpu():
     div_res = global_step / decay_steps
     if staircase:
         div_res = ops.floor(div_res)
@@ -184,7 +181,6 @@ def inverse_time_decay(learning_rate, decay_steps, decay_rate, staircase=False):
     """
     global_step = _decay_step_counter()
 
-    with init_on_cpu():
     div_res = global_step / decay_steps
     if staircase:
         div_res = ops.floor(div_res)
@@ -224,13 +220,10 @@ def polynomial_decay(learning_rate,
     """
     global_step = _decay_step_counter()
 
-    with init_on_cpu():
     if cycle:
         div_res = ops.ceil(global_step / decay_steps)
-            zero_var = tensor.fill_constant(
-                shape=[1], dtype='float32', value=0.0)
-            one_var = tensor.fill_constant(
-                shape=[1], dtype='float32', value=1.0)
+        zero_var = tensor.fill_constant(shape=[1], dtype='float32', value=0.0)
+        one_var = tensor.fill_constant(shape=[1], dtype='float32', value=1.0)
 
         with control_flow.Switch() as switch:
             with switch.case(global_step == zero_var):
@@ -277,7 +270,6 @@ def piecewise_decay(boundaries, values):
 
     global_step = _decay_step_counter()
 
-    with init_on_cpu():
     lr = tensor.create_global_var(
         shape=[1],
         value=0.0,
@@ -288,15 +280,16 @@ def piecewise_decay(boundaries, values):
     with control_flow.Switch() as switch:
         for i in range(len(boundaries)):
             boundary_val = tensor.fill_constant(
-                    shape=[1], dtype='float32', value=float(boundaries[i]))
+                shape=[1],
+                dtype='float32',
+                value=float(boundaries[i]),
+                force_cpu=True)
             value_var = tensor.fill_constant(
                 shape=[1], dtype='float32', value=float(values[i]))
             with switch.case(global_step < boundary_val):
                 tensor.assign(value_var, lr)
         last_value_var = tensor.fill_constant(
-                shape=[1],
-                dtype='float32',
-                value=float(values[len(values) - 1]))
+            shape=[1], dtype='float32', value=float(values[len(values) - 1]))
         with switch.default():
             tensor.assign(last_value_var, lr)
 

python/paddle/fluid/tests/unittests/test_learning_rate_scheduler.py

@@ -91,20 +91,21 @@ class TestLearningRateDecay(unittest.TestCase):
 
     def check_decay_with_place(self, place, python_decay_fn, fluid_decay_fn,
                                kwargs):
+        main_prog = fluid.Program()
+        startup_prog = fluid.Program()
 
+        with fluid.program_guard(main_prog, startup_prog):
             decayed_lr = fluid_decay_fn(**kwargs)
 
         place = fluid.CPUPlace()
         exe = fluid.Executor(place)
 
-        exe.run(fluid.default_startup_program())
+        exe.run(startup_prog)
 
-        fluid.memory_optimize(fluid.default_main_program())
+        fluid.memory_optimize(main_prog)
 
         for step in range(10):
-            lr_val, = exe.run(fluid.default_main_program(),
-                              feed={},
-                              fetch_list=[decayed_lr])
+            lr_val, = exe.run(main_prog, feed={}, fetch_list=[decayed_lr])
             python_decayed_lr = python_decay_fn(
                 global_step=float(step), **kwargs)
             self.assertAlmostEqual(