tensorflow分布式训练

情况一、单机单卡

单机单卡是最普通的情况，当然也是最简单的，示例代码如下：

[python] view plain copy

1. #coding=utf-8  
2. #单机单卡  
3. #对于单机单卡，可以把参数和计算都定义再gpu上，不过如果参数模型比较大，显存不足等情况，就得放在cpu上  
4. import  tensorflow as tf  
5.   
6. with tf.device('/cpu:0'):#也可以放在gpu上  
7.     w=tf.get_variable('w',(2,2),tf.float32,initializer=tf.constant_initializer(2))  
8.     b=tf.get_variable('b',(2,2),tf.float32,initializer=tf.constant_initializer(5))  
9.   
10. with tf.device('/gpu:0'):  
11.     addwb=w+b  
12.     mutwb=w*b  
13.   
14.   
15. ini=tf.initialize_all_variables()  
16. with tf.Session() as sess:  
17.     sess.run(ini)  
18.     np1,np2=sess.run([addwb,mutwb])  
19.     print np1  
20.     print np2  

情况二、单机多卡

单机多卡，只要用device直接指定设备，就可以进行训练，SGD采用各个卡的平均值，示例代码如下：

[python] view plain copy

1. #coding=utf-8  
2. #单机多卡：  
3. #一般采用共享操作定义在cpu上，然后并行操作定义在各自的gpu上，比如对于深度学习来说，我们一把把参数定义、参数梯度更新统一放在cpu上  
4. #各个gpu通过各自计算各自batch 数据的梯度值，然后统一传到cpu上，由cpu计算求取平均值，cpu更新参数。  
5. #具体的深度学习多卡训练代码，请参考：https://github.com/tensorflow/models/blob/master/inception/inception/inception_train.py  
6. import  tensorflow as tf  
7.   
8. with tf.device('/cpu:0'):  
9.     w=tf.get_variable('w',(2,2),tf.float32,initializer=tf.constant_initializer(2))  
10.     b=tf.get_variable('b',(2,2),tf.float32,initializer=tf.constant_initializer(5))  
11.   
12. with tf.device('/gpu:0'):  
13.     addwb=w+b  
14. with tf.device('/gpu:1'):  
15.     mutwb=w*b  
16.   
17.   
18. ini=tf.initialize_all_variables()  
19. with tf.Session() as sess:  
20.     sess.run(ini)  
21.     while 1:  
22.         print sess.run([addwb,mutwb])  

情况三、多机多卡

一、基本概念

Cluster、Job、task概念：三者可以简单的看成是层次关系，task可以看成每台机器上的一个进程，多个task组成job；job又有：ps、worker两种，分别用于参数服务、计算服务，组成cluster。

二、同步SGD与异步SGD

1、所谓的同步更新指的是：各个用于并行计算的电脑，计算完各自的batch 后，求取梯度值，把梯度值统一送到ps服务机器中，由ps服务机器求取梯度平均值，更新ps服务器上的参数。

如下图所示，可以看成有四台电脑，第一台电脑用于存储参数、共享参数、共享计算，可以简单的理解成内存、计算共享专用的区域，也就是ps job；另外三台电脑用于并行计算的，也就是worker task。

这种计算方法存在的缺陷是：每一轮的梯度更新，都要等到A、B、C三台电脑都计算完毕后，才能更新参数，也就是迭代更新速度取决与A、B、C三台中，最慢的那一台电脑，所以采用同步更新的方法，建议A、B、C三台的计算能力都不想。

2、所谓的异步更新指的是：ps服务器收到只要收到一台机器的梯度值，就直接进行参数更新，无需等待其它机器。这种迭代方法比较不稳定，收敛曲线震动比较厉害，因为当A机器计算完更新了ps中的参数，可能B机器还是在用上一次迭代的旧版参数值。

三、代码编写

1、定义集群

比如假设上面的图所示，我们有四台电脑，四台电脑的名字假设为：A、B、C、D，那么集群可以定义如下

[python] view plain copy

1. #coding=utf-8  
2. #多台机器，每台机器有一个显卡、或者多个显卡，这种训练叫做分布式训练  
3. import  tensorflow as tf  
4. #现在假设我们有A、B、C、D四台机器，首先需要在各台机器上写一份代码，并跑起来，各机器上的代码内容大部分相同  
5. # ，除了开始定义的时候，需要各自指定该台机器的task之外。以机器A为例子，A机器上的代码如下：  
6. cluster=tf.train.ClusterSpec({  
7.     "worker": [  
8.         "A_IP:2222",#格式 IP地址：端口号，第一台机器A的IP地址 ,在代码中需要用这台机器计算的时候，就要定义：/job:worker/task:0  
9.         "B_IP:1234"#第二台机器的IP地址 /job:worker/task:1  
10.         "C_IP:2222"#第三台机器的IP地址 /job:worker/task:2  
11.     ],  
12.     "ps": [  
13.         "D_IP:2222",#第四台机器的IP地址 对应到代码块：/job:ps/task:0  
14.     ]})  

然后我们需要写四分代码，这四分代码文件大部分相同，但是有几行代码是各不相同的。

2、在各台机器上，定义server

比如A机器上的代码server要定义如下：

[python] view plain copy

1. server=tf.train.Server(cluster,job_name='worker',task_index=0)#找到‘worker’名字下的，task0，也就是机器A  

3、在代码中，指定device

[python] view plain copy

1. with tf.device('/job:ps/task:0'):#参数定义在机器D上  
2.     w=tf.get_variable('w',(2,2),tf.float32,initializer=tf.constant_initializer(2))  
3.     b=tf.get_variable('b',(2,2),tf.float32,initializer=tf.constant_initializer(5))  
4.   
5. with tf.device('/job:worker/task:0/cpu:0'):#在机器A cpu上运行  
6.     addwb=w+b  
7. with tf.device('/job:worker/task:1/cpu:0'):#在机器B cpu上运行  
8.     mutwb=w*b  
9. with tf.device('/job:worker/task:2/cpu:0'):#在机器C cpu上运行  
10.     divwb=w/b  

在深度学习训练中，一般图的计算，对于每个worker task来说，都是相同的，所以我们会把所有图计算、变量定义等代码，都写到下面这个语句下：

[python] view plain copy

1. with tf.device(tf.train.replica_device_setter(worker_device='/job:worker/task:indexi',cluster=cluster)):  

函数replica_deviec_setter会自动把变量参数定义部分定义到ps服务中(如果ps有多个任务，那么自动分配)。下面举个例子，假设现在有两台机器A、B，A用于计算服务，B用于参数服务，那么代码如下：

[python] view plain copy

1. #coding=utf-8  
2. #上面是因为worker计算内容各不相同，不过再深度学习中，一般每个worker的计算内容是一样的，  
3. # 以为都是计算神经网络的每个batch 前向传导，所以一般代码是重用的  
4. import  tensorflow as tf  
5. #现在假设我们有A、B台机器，首先需要在各台机器上写一份代码，并跑起来，各机器上的代码内容大部分相同  
6. # ，除了开始定义的时候，需要各自指定该台机器的task之外。以机器A为例子，A机器上的代码如下：  
7. cluster=tf.train.ClusterSpec({  
8.     "worker": [  
9.         "192.168.11.105:1234",#格式 IP地址：端口号，第一台机器A的IP地址 ,在代码中需要用这台机器计算的时候，就要定义：/job:worker/task:0  
10.     ],  
11.     "ps": [  
12.         "192.168.11.130:2223"#第四台机器的IP地址 对应到代码块：/job:ps/task:0  
13.     ]})  
14.   
15. #不同的机器，下面这一行代码各不相同，server可以根据job_name、task_index两个参数，查找到集群cluster中对应的机器  
16.   
17. isps=False  
18. if isps:  
19.     server=tf.train.Server(cluster,job_name='ps',task_index=0)#找到‘worker’名字下的，task0，也就是机器A  
20.     server.join()  
21. else:  
22.     server=tf.train.Server(cluster,job_name='worker',task_index=0)#找到‘worker’名字下的，task0，也就是机器A  
23.     with tf.device(tf.train.replica_device_setter(worker_device='/job:worker/task:0',cluster=cluster)):  
24.         w=tf.get_variable('w',(2,2),tf.float32,initializer=tf.constant_initializer(2))  
25.         b=tf.get_variable('b',(2,2),tf.float32,initializer=tf.constant_initializer(5))  
26.         addwb=w+b  
27.         mutwb=w*b  
28.         divwb=w/b  
29.   
30. saver = tf.train.Saver()  
31. summary_op = tf.merge_all_summaries()  
32. init_op = tf.initialize_all_variables()  
33. sv = tf.train.Supervisor(init_op=init_op, summary_op=summary_op, saver=saver)  
34. with sv.managed_session(server.target) as sess:  
35.     while 1:  
36.         print sess.run([addwb,mutwb,divwb])  

把该代码在机器A上运行，你会发现，程序会进入等候状态，等候用于ps参数服务的机器启动，才会运行。因此接着我们在机器B上运行如下代码：

[python] view plain copy

1. #coding=utf-8  
2. #上面是因为worker计算内容各不相同，不过再深度学习中，一般每个worker的计算内容是一样的，  
3. # 以为都是计算神经网络的每个batch 前向传导，所以一般代码是重用的  
4. #coding=utf-8  
5. #多台机器，每台机器有一个显卡、或者多个显卡，这种训练叫做分布式训练  
6. import  tensorflow as tf  
7. #现在假设我们有A、B、C、D四台机器，首先需要在各台机器上写一份代码，并跑起来，各机器上的代码内容大部分相同  
8. # ，除了开始定义的时候，需要各自指定该台机器的task之外。以机器A为例子，A机器上的代码如下：  
9. cluster=tf.train.ClusterSpec({  
10.     "worker": [  
11.         "192.168.11.105:1234",#格式 IP地址：端口号，第一台机器A的IP地址 ,在代码中需要用这台机器计算的时候，就要定义：/job:worker/task:0  
12.     ],  
13.     "ps": [  
14.         "192.168.11.130:2223"#第四台机器的IP地址 对应到代码块：/job:ps/task:0  
15.     ]})  
16.   
17. #不同的机器，下面这一行代码各不相同，server可以根据job_name、task_index两个参数，查找到集群cluster中对应的机器  
18.   
19. isps=True  
20. if isps:  
21.     server=tf.train.Server(cluster,job_name='ps',task_index=0)#找到‘worker’名字下的，task0，也就是机器A  
22.     server.join()  
23. else:  
24.     server=tf.train.Server(cluster,job_name='worker',task_index=0)#找到‘worker’名字下的，task0，也就是机器A  
25.     with tf.device(tf.train.replica_device_setter(worker_device='/job:worker/task:0',cluster=cluster)):  
26.         w=tf.get_variable('w',(2,2),tf.float32,initializer=tf.constant_initializer(2))  
27.         b=tf.get_variable('b',(2,2),tf.float32,initializer=tf.constant_initializer(5))  
28.         addwb=w+b  
29.         mutwb=w*b  
30.         divwb=w/b  
31.   
32. saver = tf.train.Saver()  
33. summary_op = tf.merge_all_summaries()  
34. init_op = tf.initialize_all_variables()  
35. sv = tf.train.Supervisor(init_op=init_op, summary_op=summary_op, saver=saver)  
36. with sv.managed_session(server.target) as sess:  
37.   
38.     while 1:  
39.         print sess.run([addwb,mutwb,divwb])  

分布式训练需要熟悉的函数：

• tf.train.Server
• tf.train.Supervisor
• tf.train.SessionManager
• tf.train.ClusterSpec
• tf.train.replica_device_setter
• tf.train.MonitoredTrainingSession
• tf.train.MonitoredSession
• tf.train.SingularMonitoredSession
• tf.train.Scaffold
• tf.train.SessionCreator
• tf.train.ChiefSessionCreator
• tf.train.WorkerSessionCreator

参考文献：

https://www.tensorflow.org/versions/master/how_tos/distributed/index.html