深度学习分布式训练实战（二）——TF

本篇博客主要介绍TF的分布式训练，重点从代码层面进行讲解。理论部分可以参考深度学习分布式训练实战（一)

TF的分布式实现方式

TF的分布式有两种实现方式，一种是图内分布式（In-graph replication）；一种是图间分布式(Between-graph replication)。图内分布式中，计算图只有一个，需要一个中心节点分配计算任务并更新参数，由于中心节点的存在，中心节点容易成为瓶颈。图间分布式中，计算图有多个，但是不同计算图的相同变量通过tf.train.replica_device_setter函数放到同一个服务器上，这种情况下，各个计算图相互独立（参数只有一份，计算图有多个），并行度更高，适合异步更新，同步更新下相对麻烦，不过TF给了接口tf.train.SyncReplicasOptimizer函数来帮助实现参数的同步更新，所以图间分布式应用相对广泛一些。
关于数据并行，模型并行可以参考深度学习分布式训练实战（一)
大部分情况下，我们使用图间分布式，图内分布式一般只会在模型太大的情况下使用。对于图间分布式，其基于gRPC通信框架，模型参数只有一份，计算图有多份，一个master负责创建主session，多个worker执行计算图任务。模型训练过程中，每个计算图计算出各自梯度，然后对参数进行更新。更新方式有两种：同步更新，异步更新。

分布式TF中，TF需要建立一个集群，然后在集群中建立两个job，一个是ps job，负责参数初始化，参数更新，一个job下面可以有多个task（有多个task，说明有多台机器，或者GPU负责参数初始化，更新）。一个是woker job，负责计算图的运算，计算梯度，一个worker job下面也可以有很多个task（有多个task，说明有多台机器，或者GPU负责运行计算图）。

参数异步更新的分布式训练

参数同步更新基本上和这里写的差不多TensorFlow分布式部署
。只不过为了方便在本机上调试，所以改了一点点。（自己的笔记本没有GPU），介绍下面几个重点的语句：
tf.train.ClusterSpec()：创建一个集群对象
tf.train.Server()：在这个集群上面创建一个服务器，根据实际情况，可以是参数服务器，也可以是计算服务器
tf.train.Supervisor()：创建一个监视器，就是用来监控训练过程的，个人感觉主要就是方便恢复模型训练，其logdir参数为训练日志目录，如果里面有模型，则直接恢复训练。所以如果想重新训练，需要删除这个目录。
sv.managed_session()：启动Session，相比于其他启动Session的方式，多了一些功能。可以参考TensorFlow 中三种启动图用法

具体代码如下：

# tensorflow distribute train by asynchronously update 

import tensorflow as tf
import numpy as np

tf.app.flags.DEFINE_string("ps_hosts", "", "Comma-separated list of hostname:port pairs")
tf.app.flags.DEFINE_string("worker_hosts", "", "Comma-separated list of hostname:port pairs")
tf.app.flags.DEFINE_string("job_name", "", "one of ps or worker")
tf.app.flags.DEFINE_integer("task_index", 0, "0, 1, 2...")

FLAGS = tf.app.flags.FLAGS

def main(_):
	ps_hosts = FLAGS.ps_hosts.split(",")
	worker_hosts = FLAGS.worker_hosts.split(",")

	# Create a cluster from the parameter server and worker server
	cluster = tf.train.ClusterSpec({"ps":ps_hosts, "worker":worker_hosts})

	# Create and start a server for the local task
	server = tf.train.Server(cluster, job_name = FLAGS.job_name, task_index=FLAGS.task_index)
    # 如果是参数服务器，则直接阻塞，等待计算服务器下达参数初始化，参数更新命令就可以了。
    # 不过“下达命令”这个是TF内部实现的，没有显式实现
	if FLAGS.job_name == "ps":
		server.join() 
	elif FLAGS.job_name == "worker":
		# Assigns ops to the local worker by default
		with tf.device(tf.train.replica_device_setter(worker_device="/job:worker/task:%d" % FLAGS.task_index, cluster=cluster)):
			train_X = np.linspace(-1.0, 1.0, 100)
			train_Y = 2.0 * train_X + np.random.randn(*train_X.shape) * 0.33 + 10.0
			X = tf.placeholder("float")
			Y = tf.placeholder("float")

			w = tf.Variable(0.0, name="weight")
			b = tf.Variable(0.0, name="bias")
			loss = tf.square(Y - tf.multiply(X, w) - b)

			global_step = tf.Variable(0)
			train_op = tf.train.AdagradOptimizer(0.01).minimize(loss, global_step=global_step)

			saver = tf.train.Saver()
			summary_op = tf.summary.merge_all()

			init_op = tf.global_variables_initializer()

			# Create a "supervisor", which oversees the training process.
			sv = tf.train.Supervisor(is_chief=(FLAGS.task_index==0),
				logdir="~/Downloads/log/",
				init_op=init_op,
				summary_op = summary_op,
				saver=saver,
				global_step=global_step,
				save_model_secs=600)

			# The supervisor takes care of session initialization, retoring from a
			# checkpoint, and closing when done or an error occurs.
			with sv.managed_session(server.target) as sess:
				step = 0
				while step < 1000000:
					# Run a training step asynchronously
					for (x, y) in zip(train_X, train_Y):
						_, step =sess.run([train_op, global_step], feed_dict={X:x, Y:y})
					loss_value = sess.run(loss, feed_dict={X:x, Y:y})
					print("Step: {}, loss: {}".format(step, loss_value))

			# Ask for all the services to stop
			sv.stop()

if __name__=="__main__":
	tf.app.run()

打开三个终端，分别输入以下三个命令，就可以看到训练结果了：

CUDA_VISIBLE_DEVICES='' python AsynDis.py --ps_hosts=localhost:2222 --worker_hosts=localhost:2224,localhost:2225 --job_name=ps --task_index=0
CUDA_VISIBLE_DEVICES='' python AsynDis.py --ps_hosts=localhost:2222 --worker_hosts=localhost:2224,localhost:2225 --job_name=worker --task_index=0
CUDA_VISIBLE_DEVICES='' python AsynDis.py --ps_hosts=localhost:2222 --worker_hosts=localhost:2224,localhost:2225 --job_name=worker --task_index=1

这里有一点要注意，控制计算图使用哪一块GPU是通过命令行设置CUDA_VISIBLE_DEVICES来实现的，而不是写死在代码里面的。
此外，还有一点不方便的地方，如果有很多台机器，则需要把多份这份代码拷贝多次，在每台机器上分别运行上述命令才可以，还是不太方便的。

参数同步更新的分布式训练

同步更新稍微麻烦了点，需要加几行代码（重点参考了《Tensorflow实战》一书），改动部分已经标明，代码如下：

# tensorflow distribute train by synchronously update 

import tensorflow as tf
import numpy as np

tf.app.flags.DEFINE_string("ps_hosts", "", "Comma-separated list of hostname:port pairs")
tf.app.flags.DEFINE_string("worker_hosts", "", "Comma-separated list of hostname:port pairs")
tf.app.flags.DEFINE_string("job_name", "", "one of ps or worker")
tf.app.flags.DEFINE_integer("task_index", 0, "0, 1, 2...")

FLAGS = tf.app.flags.FLAGS
def main(_):
	ps_hosts = FLAGS.ps_hosts.split(",")
	worker_hosts = FLAGS.worker_hosts.split(",")
	n_works = len(worker_hosts)
	# Create a cluster from the parameter server and worker server
	cluster = tf.train.ClusterSpec({"ps":ps_hosts, "worker":worker_hosts})

	# Create and start a server for the local task
	server = tf.train.Server(cluster, job_name = FLAGS.job_name, task_index=FLAGS.task_index)

	if FLAGS.job_name == "ps":
		server.join()
	elif FLAGS.job_name == "worker":
		# Assigns ops to the local worker by default
		with tf.device(tf.train.replica_device_setter(worker_device="/job:worker/task:%d" % FLAGS.task_index, cluster=cluster)):
			train_X = np.linspace(-1.0, 1.0, 100)
			train_Y = 2.0 * train_X + np.random.randn(*train_X.shape) * 0.33 + 10.0
			X = tf.placeholder("float")
			Y = tf.placeholder("float")

			w = tf.Variable(0.0, name="weight")
			b = tf.Variable(0.0, name="bias")
			loss = tf.square(Y - tf.multiply(X, w) - b)

			global_step = tf.Variable(0)

			# for Syncmously updata
			# 同步更新模式下，需要等待所有计算图计算出梯度，然后梯度求平均，tf.train.SyncReplicasOptimizer实现了这种封装
			opt = tf.train.SyncReplicasOptimizer(
				tf.train.AdagradOptimizer(0.01),
				replicas_to_aggregate=n_works,
				total_num_replicas=n_works,
				)
			train_op = opt.minimize(loss, global_step=global_step)
			saver = tf.train.Saver()
			summary_op = tf.summary.merge_all()

			init_op = tf.global_variables_initializer()

			# for Syncmously updata
			# 同步模式下，主计算服务器需要协调不同计算服务器计算得到的梯度，并更新参数。
			if FLAGS.task_index==0:
			    # 定义协调不同计算服务器的队列，并定义初始化操作
				chief_queue_runner = opt.get_chief_queue_runner()
				init_tokens_op = opt.get_init_tokens_op(0)

			# Create a "supervisor", which oversees the training process.
			sv = tf.train.Supervisor(is_chief=(FLAGS.task_index==0),
				logdir="~/Downloads/log/",
				init_op=init_op,
				summary_op = summary_op,
				saver=saver,
				global_step=global_step,
				save_model_secs=600)



			# The supervisor takes care of session initialization, retoring from a
			# checkpoint, and closing when done or an error occurs.

			# for Syncmously updata. 
			# prepare_or_wait_for_session used by sync. It will wait until main node ok and parameter init over!
			# for Syncmously updata. 
			# 这里用的是prepare_or_wait_for_session。
			# 相比于同步更新的managed_session：只要某个计算服务器参数初始化完毕就可以开始，
			# prepare_or_wait_for_session：等待所有计算服务器参数初始化完毕（参数只有一份，后续的计算服务器应该不需要初始化了？只需要和参数服务器建立一个关系？），主节点协调工作完毕后，开始。
			with sv.prepare_or_wait_for_session(server.target) as sess:
				# for Syncmously updata
				if FLAGS.task_index==0:
				    # 开始训练之前，主计算服务器需要启动协调同步更新的队列，并执行初始化操作
					sv.start_queue_runners(sess, [chief_queue_runner])
					sess.run(init_tokens_op)

				step = 0
				while step < 100000:
					# Run a training step asynchronously
					for (x, y) in zip(train_X, train_Y):
						_, step =sess.run([train_op, global_step], feed_dict={X:x, Y:y})
					loss_value = sess.run(loss, feed_dict={X:x, Y:y})
					print("Step: {}, loss: {}".format(step, loss_value))

			# Ask for all the services to stop
			sv.stop()

if __name__=="__main__":
	tf.app.run()