Tensorflow分布式训练实例详解

一、Tensorflow分布式训练策略

模型并行 in-graph replication：将模型部署到很多设备上运行，比如一个或多个机器的GPUs， 不同 GPU 负责网络模型的不同部分。 数据并行 between-graph replication：每个Worker拥有模型的完整副本但分配到不同的数据，各自单独地训练，然后将计算结果合并。这是目前主流的做法。   同步更新： 等到所有 GPU 的梯度计算完成，再统一根据梯度均值计算新参数，然后所有 GPU 同步新值后，才进行下一轮计算，loss下降比较稳定，但需要等待最慢的计算节点。 异步更新：所有 GPU 分别计算，分别更新参数，计算资源能充分利用，但loss下降不稳定，存在梯度失效问题。     1. gRPC (google remote procedure call) TensorFlow分布式并行基于gRPC通信框架（谷歌远程过程调用），其中包括一个master创建Session，还有多个worker负责执行计算图中的任务。 即：一个Cluster可以切分多个Job，一个Job指一类特定的任务，每个Job包含多个Task，比如parameter server(ps)、worker，在大多数情况下,一个机器上只运行一个Task. Cluster是Job的集合，Job是Task的集合。 Parameter Server是管理参数的存储和更新工作. Worker是来运行ops. 在每个迭代过程，worker从parameter sever中获得参数，然后将计算的梯度返回给parameter server，parameter server聚合从worker传回的梯度，然后更新参数，并将新的参数广播给worker。 2.  Ring AllReduce架构     PS架构中，当worker数量较多时，ps节点的网络带宽将成为系统的瓶颈。     Ring AllReduce架构中各个设备都是worker，没有中心节点来聚合所有worker计算的梯度。所有 device 在一个逻辑环路中，每个 device 从上  行的device 接收数据，并向下行的 deivce 发送数据。

 

二、TensorFlow Estimator API

首先定义一个由参与分布式计算的机器组成的集群，集群中一般有多个worker，需要指定其中一个worker为主节点（cheif），chief节点会执行一些额外的工作，比如模型导出之类的。在PS分布式架构环境中，还需要定义ps节点。如下： cluster = {'chief': ['host0:2222'], 'ps': ['host1:2222', 'host2:2222'], 'worker': ['host3:2222', 'host4:2222', 'host5:2222']}     设置好TF_CONFIG环境变量： # Example of non-chief node: os.environ['TF_CONFIG'] = json.dumps( {'cluster': cluster, 'task': {'type': 'worker', 'index': 1}}) # Example of chief node: os.environ['TF_CONFIG'] = json.dumps( {'cluster': cluster, 'task': {'type': 'chief', 'index': 0}}) # Example of evaluator node (evaluator is not part of training cluster) os.environ['TF_CONFIG'] = json.dumps( {'cluster': cluster, 'task': {'type': 'evaluator', 'index': 0}})

 

三、TensorFlow分布式MNIST手写字体识别实例

项目地址： https://github.com/TracyMcgrady6/Distribute_MNIST 这个项目的分布式架构是用tf.train.ClusterSpec 和tf.train.Server构建的，与TensorFlow Estimator API的定义方式有所区别 # 设置job name参数 flags.DEFINE_string('job_name', None, 'job name: worker or ps') # 设置任务的索引 flags.DEFINE_integer('task_index', None, 'Index of task within the job') # 参数服务器节点 flags.DEFINE_string('ps_hosts', 'localhost:22') # 两个worker节点 flags.DEFINE_string('worker_hosts', 'localhost:23,localhost:24') # 定义任务集合 cluster = tf.train.ClusterSpec({'ps': FLAGS.ps_hosts.split(','), 'worker': FLAGS.worker_hosts.split(',')}) # tf的Sever及Session server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index) sv = tf.train.Supervisor(is_chief=is_chief, logdir='logs', init_op=init_op, recovery_wait_secs=1, global_step=global_step) sess = sv.prepare_or_wait_for_session(server.target) ps 节点执行： python distributed.py --job_name=ps --task_index=0 worker1 节点执行： python distributed.py --job_name=worker --task_index=0 worker2 节点执行： python distributed.py --job_name=worker --task_index=1