基于TensorFlow的多机多卡分布式训练搭建方法

最新推荐文章于 2024-08-08 18:13:10 发布

BmwGaara

最新推荐文章于 2024-08-08 18:13:10 发布

阅读量6.6k

点赞数 1

文章标签： tensorflow 分布式

本文链接：https://blog.csdn.net/BmwGaara/article/details/94021563

版权

1.概述

TensorFlow分布式是基于GRPC库实现的高性能集群训练框架，能有效的利用多机多卡资源，将大型的模型或者代码拆分到各个节点分别完成，从而实现高速的模型训练。

如下图所示，tensorflow的分布式集群中存在的节点主要有两种：ps节点和worker节点，ps节点是用于保存和计算训练参数的节点；worker节点是用于训练的节点。由于ps和worker节点都有可能存在多个，因此ps和worker会各存在一个master节点来统筹计算结果。

TensorFlow的分布式场景主要分成单机多卡和多机多卡两种场景，在这两种场景下，使用的方法和代码模型的结构是完全不同的，后续会详细说明。

TensorFlow的分布式是一种并行训练的训练结构，目前有两种并行方式：

1）图并行

图并行是指将一个大型的模型拆分成若干独立的小结构分别部署到不同的节点上，这样的优点是可以在不同的节点同时运行一张图中不同的部分，效率较高。但是其缺点也是非常明显的：如果两个部分存在复杂的逻辑先后关系，则在部署集群时是必须要考虑节点间的逻辑关系，从而对于部署集群造成了很大的技术困难，如果图较为简单，这种方式并不推荐。

2）数据并行

数据并行就是在各个节点上执行的图是一样的，但是跑的数据会存在差异，这种方式原理简单，利于集群的构建，是本篇描述的主要方法。

在数据并行中，还存在两种模式：图内复制和图间复制。图内复制和图间复制指的是分布式训练时各节点读取数据的方式。

图内复制是指数据都在一个节点上，这样的好处是配置简单，其他多机多GPU的计算节点，只要起个连接操作，暴露一个网络接口，等在那里接受任务就好了。但是这样的坏处是训练数据的分发在一个节点上，要把训练数据分发到不同的机器上，严重影响并发训练速度。在大数据训练的情况下，不推荐使用这种模式。

图间复制是指训练的参数保存在参数服务器，数据不用分发，数据分片的保存在各个计算节点，各个计算节点自己算自己的，算完了之后，把要更新的参数告诉参数服务器，参数服务器更新参数。这种模式的优点是不用训练数据的分发了，尤其是在数据量在TB级的时候，节省了大量的时间，所以大数据深度学习还是推荐使用图间模式。

2.多机多卡环境搭建步骤

2.1环境说明

在本篇的叙述中，我们采用了3台服务器来构建tensorflow多机多卡分布式训练的集群，其中：

ps节点：********

worker节点：********，********（前者为主节点）

2.2核心步骤说明

2.2.1参数获取

由于tensorflow的分布式服务功能由tensorflow内部支持并完成，因此ps和worker的节点代码是可以合并在同一个文件中，因此可以通过输入参数的方式来识别当前的节点类型。

flags = tf.flags
flags.DEFINE_string("job_name", "", "Either 'ps' or 'worker'")
flags.DEFINE_integer("task_index", 0, "Index of task within the job")

其中job_name用于区分ps和worker节点类型，task_index用于区分是否是主节点。

这样就可以通过类似于：

python xxx.py --job_name="ps" --task_index=0

这样的命令行来启动当前的分布式节点。

2.2.2构建集群对象

在启动各个节点服务前，tensorflow的分布式服务要求你提前告知服务你的集群到底有多少节点，并且ps和worker节点分别有哪些？

cluster = tf.train.ClusterSpec({"ps": ps_servers, "worker": workers})

在这里，ps_servers和workers分别式是两个host:port形式的list。

注意：这里节点地址列表中节点的个数会与你的index一一对应，例如你的workers里有两个服务器ip，那么task_index不能设置为2，只能是0或者1。

2.2.3构建服务对象

构建这样一个服务对象，目的在于通过传入集群信息cluster、job_name和task_index等信息将集群内各个节点通过web链接起来。

server = tf.train.Server(
    server_or_cluster_def=cluster,
    job_name=FLAGS.job_name,
    task_index=FLAGS.task_index)

2.2.4启动ps节点

ps节点的启动非常简单，简单到几乎不需要人为的去进行设置：

if FLAGS.job_name == "ps":
    server.join()

2.2.5启动worker节点

启动worker节点就要复杂很多。

首先，假设我们判断其为worker节点：

1）设置资源使用参数：

with tf.device(tf.train.replica_device_setter(
        worker_device="/job:worker/task:%d" % FLAGS.task_index,
        cluster=cluster)):
        神经网络定义
        ...

sess_config = tf.ConfigProto(allow_soft_placement=True,
                             log_device_placement=False,
                             gpu_options=tf.GPUOptions(allow_growth=True))

注意：这里可以使用指定当前节点某个甚至是某些GPU来训练，例如：worker_device="/job:worker/task:%d/gpu:0" % FLAGS.task_index。但是这种形式如果对GPU显存没有绝对的信心，是很容易造成GPU内存溢出的情况，从而导致程序报错退出训练。

这里并没有指定某个具体的GPU进行训练，配合上sess_confg中的gpu_options这个参数，tensorflow会自动调用这个节点的计算资源进行训练，从而避免内存溢出的情况。如下图：

这里是某个节点某一时刻的GPU使用截图，可以看到，tensorflow动态分配GPU的方式进行分配资源，空闲的GPU都会得到调用。

2）手动使用GPU资源

那么当一个集群可能同时接受并训练多个模型时，如何手动使用各个节点的GPU资源呢？这里提供一个简单的思路和方法。

首先通过预知训练的规模和GPU的数量，限制在该任务下，该节点可以使用的GPU有哪些：

os.environ['CUDA_VISIBLE_DEVICES'] = gpu_no

再限定这些核心可以占用的内存上限：

config=tf.ConfigProto(gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.5))

以这样配置组合形式来手动设置当前节点对这个任务的支持力度。

3）使用supervisor

supervisor是tensorflow提供的一个集记录events文件、ckpt文件、初始化会话等多个功能的对象，在分布式训练中使用该模块可以有效的降低代码量提升开发效率。

sv = tf.train.Supervisor(is_chief=(FLAGS.task_index == 0),
logdir=os.path.join(data_path, 'summary'),
init_op=init_op,
summary_op=summary_op,
global_step=global_step,
saver=None)

这里的is_chief表示当前节点是否是主节点，因为训练的结果只会保存在主节点上。logdir指的是你训练结果的保存目录，如果不设置saver和summary_op，那么所有的文件都会保存在logdir里。

但是为了将结果区分开，本人将supervisor内部的saver置空，再指定save路径，从而将events文件和ckpt文件分离开。

最后我们看下两个worker工作的截图，两个worker的训练结果会交替上升，不断迭代更新结果：