基于UAI Train平台的分布式训练

 前言

在大型数据集上进行训练的现代神经网络架构可以跨广泛的多种领域获取可观的结果，领域涵盖从图像识别、自然语言处理，到欺诈检测和推荐系统等各个方面。但是训练这些神经网络模型需要大量浮点计算能力。尽管近些年来 GPU 硬件算力和训练方法上均取得了重大的进步，但事实是在单一机器上，网络训练所需要的时间仍然长得不切实际。因此需要借助分布式GPU环境来提升神经网络训练系统的浮点计算能力。

TensorFlow分布式训练

TensorFlow 采用了数据流范式， 使用节点和边的有向图来表示计算。TensorFlow 需要用户静态声明这种符号计算图，并对该图使用复写和分区（rewrite & partitioning）将其分配到机器上进行分布式执行。

TensorFlow 中的分布式机器学习训练使用了如图所示的参数服务器方法 。

Cluster、Job、Task

关于TensorFlow的分布式训练，主要概念包括Cluster、Job、Task，其关联关系如下：

• TensorFlow分布式Cluster由多个Task组成，每个Task对应一个tf.train.Server实例，作为Cluster的一个单独节点；
•  多个相同作用的Task可以被划分为一个Job，在分布式深度学习框架中,我们一般把Job划分为Parameter Server和Worker，Parameter Job是管理参数的存储和更新工作，而Worker Job运行OPs，作为计算节点只执行计算密集型的Graph计算；
•  Cluster中的Task会相对进行通信，以便进行状态同步、参数更新等操作，如果参数的数量过大，一台机器处理不了，这就要需要多个Task。

TensorFlow分布式计算模式

In-graph 模式

In-graph模式，将模型的计算图的不同部分放在不同的机器上执行。 把计算从单机多GPU扩展到了多机多GPU， 不过数据分发还是在一个节点。 这样配置简单， 多机多GPU的计算节点只需进行join操作， 对外提供一个网络接口来接受任务。训练数据的分发依然在一个节点上， 把训练数据分发到不同的机器上， 将会影响并发训练速度。在大数据训练的情况下， 不推荐使用这种模式。

Between-graph 模式

Between-graph模式下，数据并行，每台机器使用完全相同的计算图。训练的参数保存在参数服务器，数据不用分发， 数据分布在在各个计算节点， 各个计算节点自行计算， 把要更新的参数通知参数服务器进行更新。这种模式不需要再练数据的分发， 数据量在TB级时可以节省大量的时间，目前主流的分布式训练模式以 Between-graph为主。

参数更新方式

同步更新

各个用于并行计算的节点，计算完各自的batch 后，求取梯度值，把梯度值统一送到PS参数服务机器中，并等待PS更新模型参数。PS参数服务器在收集到到一定数量计算节点的梯度后，求取梯度平均值，更新PS参数服务器上的参数，同时将参数推送到各个worker节点。

异步更新

PS参数服务器收到只要收到一台机器的梯度值，就直接进行参数更新，无需等待其它机器。这种迭代方法比较不稳定，因为当A机器计算完更新了PS参数服务器中的参数，可能B机器还是在用上一次迭代的旧版参数值。

分布式训练步骤

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. 如果当前节点是Parameter Server，则调用server.join()无休止等待；如果是Worker，则执行下一步

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

4. 则构建要训练的模型

# build tensorflow graph model 

5. 创建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 

UAI Train 分布式训练部署

UCloud AI 训练服务（UCloud AI Train）是面向AI训练任务的大规模分布式计算平台，基于高性能GPU计算节点提供一站式托管AI训练任务服务，用户在提交AI训练任务后无需担心计算节点调度、训练环境准备、数据上传下载以及容灾等问题。

目前UAI Train平台支持TensorFlow 和 MXNet 框架的分布式训练。需要将PS代码和Worker代码实现在同一个代码入口中，在执行过程中，PS 和 Worker 将使用相同的Docker容器镜像和相同的python代码入口进行执行，系统将自动生成PS和Worker的env环境参数。TensorFlow 分布式训练采用PS-Worker的分布式格式，并提供python的接口运行分布式训练。

UCloud AI Train分布式训练采用Parameter Server和Worker Server混合部署的方法，所有计算节点均由GPU物理云主机组成。PS 仅使用CPU进行计算，Worker Server则同时使用GPU和CPU进行计算，PS 和 Worker的比例为1:1。

数据存储

分布式训练所使用的输入数据和输入数据可以是来自不同的数据源，目前UAI Train仅支持UFS作为数据的存储。

Input 数据存储

指定一个UFS网盘作为Input数据源，UAI Train平台在训练执行过程中会将对应的UFS数据映射到训练执行的Worker容器的 /data/data 目录下，系统会自动将数据映射到执行的容器中，如 ip:/xxx/data/imagenet/tf → /data/data/。

Output 数据存储

指定一个UFS网盘作为output数据源，UAI Train平台在训练执行过程中会将对应的UFS数据映射到训练执行的每一个PS容器和Worker容器的 /data/output 目录下，并以共享的方式访问同一份数据。同时，在训练过程您可以通过其他云主机实时访问训练保存的模型checkpoint。

案例研究：通过CIFAR-10进行图像识别

CIFAR-10是机器学习中常见的图像识别数据集，该数据集共有60000张彩色图像，这些图像，分为10个类，每类6000张图。这里面有50000张用于训练，另外10000用于测试。

http://groups.csail.mit.edu/vision/TinyImages/

调整训练代码

为了在UAI平台上进行训练，首先下载源代码，并对cifar10_main.py上做如下修改：

1. 添加相关参数：--data_dir, --output_dir, --work_dir, --log_dir, --num_gpus，UAI Train平台将会自动生成这些参数；

2. 在代码中增加UAI参数：使用data_dir配置输入文件夹、使用output_dir配置输出文件夹。

具体案例代码可以在https://github.com/ucloud/uai-sdk/tree/master/examples/tensorflow/train/cifar 获取

在UAI Train平台执行训练

1. 根据 https://github.com/tensorflow/models/tree/master/tutorials/image/cifar10_estimator 的说明生成CIFAR-10的tfrecords；

2. 使用UAI-SDK提供的tf_tools.py 生成CIFAR-10样例的Docker镜像；

3. 确保Docker镜像已经上传至UHub，在UAI Train平台上执行。

/data/cifar10_main.py --train-batch-size=16 

在UAI平台上的分布式训练

CIFAR-10样例代码使用tf.estimator.Estimator API，只需一个分布式环境和分布式环境配置便可直接进行分布式训练，该配置需要适用于tf.estimator.Estimator API的标准，即定义一个TF_CONFIG 配置。

TF_CONFIG = { 
    "cluster":{ 
        "master":["ip0:2222"], 
        "ps":["ip0:2223","ip1:2223"], 
        "worker":["ip1:2222"]}, 
    "task":{"type":"worker","index":0}, 
    "environment":"cloud" 
} 

UAI Train平台的分布式训练功能可以自动生成TensorFlow分布式训练的GPU集群环境，同时为每个训练节点自动生成TF_CONFIG。因此，在UAI Train平台上执行CIFAR-10的分布式训练和单机训练一样，仅需要指定input/output的UFS地址并执行如下指令即可：

/data/cifar10_main.py --train-batch-size=16 

总结

UAI Train TensorFlow的分布式训练环境实现基于TensorFlow 的分布式训练系统实现，采用默认的grpc协议进行数据交换。PS和Worker采用混合部署的方式部署，PS使用纯CPU计算，Worker使用GPU+CPU计算。

在UAI Train平台中可以非常方便的开展分布式计算，提高效率、压缩训练时间。本文中最后通过CIFAR-10 案例进行解析在UAI Train平台上进行训练所需进行的修改，并在UAI Train平台上进行分布式训练。