简介
Tensorflow API提供了Cluster、Server以及Supervisor来支持模型的分布式训练。
关于Tensorflow的分布式训练介绍可以参考Distributed Tensorflow文档。简单的概括说明如下:
- Tensorflow分布式Cluster由多个Task组成,每个Task对应一个tf.train.Server实例,作为Cluster的一个单独节点。
- 多个相同作用的Task可以被划分为一个job,例如ps job作为参数服务器只保存Tensorflow model的参数,而worker job则作为计算节点只执行计算密集型的Graph计算。
- Cluster中的Task会相对进行通信,以便进行状态同步、参数更新等操作。
Tensorflow分布式集群的所有节点执行的代码是相同的。分布式任务代码具有固定的模式:
# 第1步:命令行参数解析,获取集群的信息ps_hosts和worker_hosts,以及当前节点的角色信息job_name和task_index
# 第2步:创建当前task结点的Server
cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index)
# 第3步:如果当前节点是ps,则调用server.join()无休止等待;如果是worker,构建要训练的模型。
if FLAGS.job_name == "ps":
server.join()
elif job_name == "worker":
is_chief = task_index == 0
##build tensorflow graph model##
# 第4步:创建tf.train.Supervisor来管理模型的训练过程
# Create a "supervisor", which oversees the training process.
sv = tf.train.Supervisor(is_chief=(FLAGS.task_index == 0), logdir="/tmp/train_logs")
# The supervisor takes care of session initialization and restoring from a checkpoint.
sess = sv.prepare_or_wait_for_session(server.target)
# Loop until the supervisor shuts down
while not sv.should_stop()
# train model
读取训练文件的时候也要根据worker number对文件进行分片
tf_record_pattern = os.path.join(dir_name, '*.tfrecords*')
files = tf.gfile.Glob(tf_record_pattern)
queue_name = "file_queue"
# split input files across workers, if specified
if task_index is not None and num_workers is not None:
num_files = len(files)
files = files[task_index:num_files:num_workers]
queue_name = "file_queue_{0}".format(task_index)
print("files: {0}".format(files))
filename_queue = tf.train.string_input_producer(files, num_epochs=epoch_number, name=queue_name)
TensorflowOnSpark
TensorflowOnSpark 支持使用Spark/Hadoop集群分布式的运行Tensorflow,号称支持所有的Tensorflow操作。需要注意的是用户需要对原有的TF程序进行简单的改造,就能够运行在Spark集群之上。
基于上述代码框架稍作修改:
# main函数改为main_fun(args, ctx),由上层pyspark主程序调用
# 第1步:命令行参数解析,获取集群的信息以及当前节点的角色信息
worker_num = ctx.worker_num
job_name = ctx.job_name
task_index = ctx.task_index
cluster_spec = ctx.cluster_spec
num_workers = len(cluster_spec['worker'])
# 第2步:创建当前task结点的Server
cluster, server = TFNode.start_cluster_server(ctx, 1, args.rdma)
# 第3步:如果当前节点是ps,则调用server.join()无休止等待;如果是worker,构建要训练的模型。
if FLAGS.job_name == "ps":
server.join()
elif job_name == "worker":
is_chief = task_index == 0
##build tensorflow graph model##
# 第4步:创建tf.train.Supervisor来管理模型的训练过程
# Create a "supervisor", which oversees the training process.
sv = tf.train.Supervisor(is_chief=(FLAGS.task_index == 0), logdir="/tmp/train_logs")
# The supervisor takes care of session initialization and restoring from a checkpoint.
sess = sv.prepare_or_wait_for_session(server.target)
# Loop until the supervisor shuts down
while not sv.should_stop()
# train model
pyspark主程序:
from pyspark.context import SparkContext
from pyspark.conf import SparkConf
from tensorflowonspark import TFCluster, TFNode
import argparse
import train_module
sc = SparkContext(conf=SparkConf().setAppName("appname"))
executors = sc._conf.get("spark.executor.instances")
num_executors = int(executors) if executors is not None else 1
parser = argparse.ArgumentParser()
parser.add_argument("-e", "--epochs", help="number of epochs", type=int, default=0)
parser.add_argument("-f", "--format", help="example format: (csv|pickle|tfr)", choices=["csv", "pickle", "tfr"], default="tfr")
parser.add_argument("-i", "--input", help="HDFS path to input train file in parallelized format")
…………
…………
args = parser.parse_args()
cluster = TFCluster.run(sc, train_module.main_fun, args, args.cluster_size, num_ps, args.tensorboard, TFCluster.InputMode.TENSORFLOW, log_dir=args.model)
cluster.shutdown()
分布式TensorFlow与Spark对比
- 分布式的级别不同:TensorFlow的Tensor、Variable和Op不是分布式的,分布式执行的是subgraph. Spark的op和变量都是构建在RDD上,RDD本身是分布式的。
- 异步训练:TensorFlow支持同步和异步的分布式训练;Spark原生的API只支持同步训练
- 分布式存储:Spark在底层封装好了worker和分布式数据之间的关系;TensorFlow需要自行维护。
- Parameter Server:TensorFlow支持,Spark暂不支持。
Distributed Tensorflow文档翻译
创建一个集群
一个Tensorflow Cluster是一个tasks的集合,这些tasks会参与一个TensoFlow图的分布式执行。每一个task与一个Tensorflow server相关,server包含一个master:用于创建sessions,一个worker:用于执行图中的操作。一个cluster也可以被分解为一个或多个jobs,每个job包含一个或多个task。 为创建一个cluster,需要为cluster中的每个task创建一个TensorFlow server。每个task通常运行在不同的机器上,但你也可以运行多个tasks在同一台机器上(比如,控制不同的GPU设备)。每一个task,做如下操作:
- 创建一个tf.train.ClusterSpec,描述一个集群中的所有tasks。对每一个task应该都一样。
- 创建tf.train.Server,传递tf.train.ClusterSpec给构造函数,并通过job name和task index来区分本地的任务。
创建一个tf.train.ClusterSpec来描述集群
集群配置词典会把job name映射到网络地址列表,把这个词典传递给tf.train.ClusterSpec构造函数,例如:
在每个task内创建一个tf.train.Server实例
一个server也可以与集群中的其他server通信。一个tf.train.Server对象包含一个本地设备的集合,在它的tf.train.ClusterSpec里面保存了与其他任务连接关系的集合和一个tf.Session,它可以使用这些信息去执行一个分布式计算 。每一个服务器是一个有具体名字的工作的成员,并在工作内部有一个任务ID。一个服务器也可以与集群中的其他服务器通信。
例如,在localhost:2222和localhost:2223上执行一个包含两个server的cluster,在本地机器上的两个不同的进程中运行下面的代码段:
In task 0:
cluster = tf.train.ClusterSpec({"local": ["localhost:2222", "localhost:2223"]})
server = tf.train.Server(cluster, job_name="local", task_index=0)
In task 1:
cluster = tf.train.ClusterSpec({"local": ["localhost:2222", "localhost:2223"]})
server = tf.train.Server(cluster, job_name="local", task_index=1)
在模型中指定分布式设备
为了把操作放到特定的进程中,你同样可以使用tf.device函数,这个函数被用来指定ops运行在CPU还是GPU上,例如:
with tf.device("/job:ps/task:0"):
weights_1 = tf.Variable(...)
biases_1 = tf.Variable(...)
with tf.device("/job:ps/task:1"):
weights_2 = tf.Variable(...)
biases_2 = tf.Variable(...)
with tf.device("/job:worker/task:7"):
input, labels = ...
layer_1 = tf.nn.relu(tf.matmul(input, weights_1) + biases_1)
logits = tf.nn.relu(tf.matmul(layer_1, weights_2) + biases_2)
# ...
train_op = ...
with tf.Session("grpc://worker7.example.com:2222") as sess:
for _ in range(10000):
sess.run(train_op)
在上面的例子中,变量在ps job中的的2个tasks中创建,模型中计算密集的部分在worker job中创建。Tensorflow会将这些数据在不同的jobs间传递(前向传播时从ps到worker,应用梯度时从worker到ps)