Tensorflow——同步模式中work节点与ps节点之间的梯度交换过程

Tensorflow在分布式下的梯度更新通常采用同步更新或者异步更新，在这里就不在赘述他们的优缺点了。在同步更新下，一般会用到一下函数

	Sync_op =   tf.train.SyncReplicasOptimizer(optimizer,
									replicas_to_aggregate=len( worker_hosts),
                                    replica_id=FLAGS.task_index,
                                    total_num_replicas=len(worker_hosts),
                                    use_locking=True)`

在这篇博客中，我简单介绍一下work节点与ps节点梯度传递的过程。在Sync_op执行操作之前，需要定义一个优化器optimizer作为Sync_op的输入参数，在SyncReplicasOptimizer()函数源码中，它这样解释optimizer这个参数

 opt: The actual optimizer that will be used to compute and apply the
        gradients. Must be one of the Optimizer classes.

这说明它实际apply_gradients到Graph中的优化器，SyncReplicasOptimizer中的apply_gradients只是封装了该函数，并添加了一些同步操作，因此对于聚合梯度并求平均值的操作也在该函数中。

所以同步更新整个的流程大概就是，worker先自定义优化器，如梯度下降优化器，再用该优化器求得整个模型的梯度，

optimizer = tf.train.GradientDescentOptimizer(learning_rate)
grads_and_vars = optimizer.compute_gradients(loss)

在这里，也可以针对grads_and_vars自定义一些操作，比如剪裁之类的，然后，即可上传到ps中，同时ps会定义一个相关操作:

 train_op = Sync_op.apply_gradients(grads_and_vars,global_step=global_step)

这里的Sync_op就是之前初始化后的同步优化器，grads_and_vars则是每个worker局部获取到的梯度，源码对于global_step的解释大意为“在这些变量经过一轮的更新之后，它的值就会加1”，它可以用来检验需要更新的变量是否过时。

接下来就是具体执行梯度聚合以及求平均值的过程，该函数首先会定义三个变量 train_ops = []
，aggregated_grad = []，var_list = []用来存储op，梯度，以及对应的变量。`

 with ops.name_scope(None, self._name):
      for grad, var in grads_and_vars:
        var_list.append(var)  #对传入grads_and_vars存取相应的变量名
        with ops.device(var.device):
          # Dense gradients.
          if grad is None:
            aggregated_grad.append(None)  # pass-through.
            continue
          elif isinstance(grad, ops.Tensor):   #酱每个worker中的梯度放入这个accumulator中
          									   #对梯度值进行聚合
            grad_accum = data_flow_ops.ConditionalAccumulator(
                grad.dtype,
                shape=var.get_shape(),
                shared_name=var.name + "/grad_accum")
            train_ops.append(grad_accum.apply_grad(
                grad, local_step=self._local_step))
	            aggregated_grad.append(grad_accum.take_grad( #聚合后对梯度取平均值 副本(replicas_to_aggregate)
                self._replicas_to_aggregate))
          else:
            if not isinstance(grad, ops.IndexedSlices):
              raise ValueError("Unknown grad type!")
            grad_accum = data_flow_ops.SparseConditionalAccumulator(
                grad.dtype, shape=(), shared_name=var.name + "/grad_accum")
            train_ops.append(grad_accum.apply_indexed_slices_grad(
                grad, local_step=self._local_step))
            aggregated_grad.append(grad_accum.take_indexed_slices_grad(
                self._replicas_to_aggregate))

          self._accumulator_list.append((grad_accum, var.device))

      aggregated_grads_and_vars = zip(aggregated_grad, var_list)#将梯度值与变量合在一起

这只是基于个人的理解，并不透彻。