Fix dist error with lr decay layer

python/paddle/fluid/distribute_transpiler.py

@@ -135,6 +135,7 @@ def transpile(self,
                   optimize_ops,
                   params_grads,
                   trainer_id,
+                  lr_decay_ops=[],
                   program=None,
                   pservers="127.0.0.1:6174",
                   trainers=1,
@@ -186,6 +187,7 @@ def transpile(self,
         self.program = program
         self.trainers = trainers
         self.optimize_ops = optimize_ops
+        self.lr_decay_ops = lr_decay_ops
         # TODO(typhoonzero): currently trainer_id is fetched from cluster system
         # like Kubernetes, we should port this to use etcd later when developing
         # fluid distributed training with fault-tolerance.
@@ -338,15 +340,23 @@ def __append_optimize_op__(op, block):
             else:
                 self._append_pserver_non_opt_ops(block, op)
 
+        append_block = optimize_block
+        # append lr decay ops to the child block if exits
+        if self.lr_decay_ops:
+            for _, op in enumerate(self.lr_decay_ops):
+                self._append_pserver_non_opt_ops(append_block, op)
+
+            append_block = pserver_program.create_block(append_block.idx)
+
         # append op to the current block
-        per_opt_block = optimize_block
+        per_opt_block = append_block
         for _, opt_op in enumerate(opt_op_on_pserver):
             for _, op in enumerate(self.optimize_ops):
                 # optimizer is connected to itself
                 if ufind.is_connected(op, opt_op) and \
                     op not in global_ops:
                     __append_optimize_op__(op, per_opt_block)
-            per_opt_block = pserver_program.create_block(0)
+            per_opt_block = pserver_program.create_block(append_block.idx)
 
         # append global ops
         for glb_op in global_ops:

python/paddle/fluid/optimizer.py

@@ -24,6 +24,7 @@
 from regularizer import append_regularization_ops
 from clip import append_gradient_clip_ops, error_clip_callback
 from contextlib import contextmanager
+from distribute_transpiler import UnionFind
 
 __all__ = [
     'SGD', 'Momentum', 'Adagrad', 'Adam', 'Adamax', 'DecayedAdagrad',
@@ -172,6 +173,42 @@ def _get_accumulator(self, name, param):
                             format(name, param.name))
         return self._accumulators[name][param.name]
 
+    def _get_lr_decay_ops(self):
+        def __is_op_connected(op1, op2):
+            op1_input_names = op1.input_arg_names
+            op1_output_names = op1.output_arg_names
+
+            op2_input_names = op2.input_arg_names
+            op2_output_names = op2.output_arg_names
+
+            if set(op1_output_names) & set(op2_input_names) or \
+               set(op1_input_names) & set(op2_output_names):
+                return True
+            return False
+
+        ret_ops = []
+        if isinstance(self._learning_rate, framework.Variable):
+            output_op_idx = -1
+            global_block = framework.default_main_program().global_block()
+
+            for idx, op in enumerate(global_block.ops):
+                if self._learning_rate.name in op.output_arg_names:
+                    output_op_idx = idx
+                    break
+            sliced_ops = global_block.slice_ops(0, output_op_idx + 1)
+            ufind = UnionFind(sliced_ops)
+            for _, op1 in enumerate(sliced_ops):
+                for _, op2 in enumerate(sliced_ops):
+                    if op1 != op2 and __is_op_connected(op1, op2):
+                        ufind.union(op1, op2)
+
+            for _, op in enumerate(sliced_ops):
+                if ufind.is_connected(op, global_block.ops[output_op_idx]):
+                    ret_ops.append(op)
+            ret_ops.append(global_block.ops[output_op_idx])
+
+        return ret_ops
+
     def create_optimization_pass(self,
                                  parameters_and_grads,
                                  loss,
@@ -217,9 +254,11 @@ def create_optimization_pass(self,
             # Get custom finish ops for subclasses
             # FIXME: Need to fix this once we figure out how to handle dependencies
             self._finish_update(loss.block)
-
             end = len(global_block.ops)
-            return global_block.slice_ops(start, end)
+
+            lr_decay_ops = self._get_lr_decay_ops()
+            optimize_ops = global_block.slice_ops(start, end)
+            return lr_decay_ops, optimize_ops
 
     def minimize(self,
                  loss,
@@ -242,9 +281,9 @@ def minimize(self,
         params_grads = append_regularization_ops(params_grads,
                                                  self.regularization)
 
-        optimize_ops = self.create_optimization_pass(params_grads, loss,
-                                                     startup_program)
-        return optimize_ops, params_grads
+        lr_decay_ops, optimize_ops = self.create_optimization_pass(
+            params_grads, loss, startup_program)
+        return lr_decay_ops, optimize_ops, params_grads,
 
 
 class SGDOptimizer(Optimizer):