【Hello AI】使用AIACC-Training TensorFlow版

TensorFlow目前进行数据分布式训练的主流方式是Horovod，AIACC-Training 1.5支持使用Horovod API兼容的方式对TensorFlow分布式训练进行加速。本文为您介绍使用AIACC-Training TensorFlow版的具体操作及可能遇到的问题。

适配Horovod API

本小节介绍如何使用Horovod兼容API进行TensorFlow分布式训练的基本步骤，以下操作为原始训练代码适配到AIACC-Traninig的一般过程。

AIACC-Training for TensorFlow支持Horovod API。适配AIACC-Training的方式与Horovod一致，如果您之前是使用Horovod进行分布式训练，只需替换import模块即可。替换内容后的内容如下所示：

import perseus.tensorflow.horovod as hvd

如果您的训练代码是非分布式代码，可以参考以下操作步骤将训练代码升级为Horovod接口的分布式训练代码。

1. 在main函数的开头部分，执行如下命令，初始化Perseus Horovod模块。

说明：请务必在使用其他Perseus API之前进行调用。

hvd.init()

2. 将训练数据集切分为子数据集。

每个进程需要使用不同的数据，来达到数据并行训练的目的，支持手动或自动进行数据切分。TensorFlow为tf.data.Dataset类提供了自动切分数据的shard()接口，可结合size()、local_rank()或rank()函数进行自定义的数据切分，命令如下：

# 每个机器使用同一份数据
dataset = tf.dataset.shard(hvd.size(), hvd.rank())
# 已将数据切分到每个机器上，每个机器使用数据的不同子集
dataset = tf.dataset.shard(hvd.size(), hvd.local_rank())

函数说明如下：

• size()：返回所有参与训练的GPU卡数，例如双机，每个机器8张卡，则返回值为16。您也可以用该参数来调整学习率、以及checkpoint保存的batch number等。
• local_rank()：返回当前GPU在本机器上的编号，例如两台8卡机器，则每台机器上的卡的local_rank()范围是0~7。您也可以用该参数来绑定进程到相应的GPU卡。
• rank()：返回当前GPU在所有机器及GPU中的编号，例如四台8卡机器，rank()的范围是0~31。

重要

• 为保证算法的准确性和性能的稳定性，请将shard()函数置于数据集的shuffle()和repeat()操作之前。
• 请勿对evaluation或test数据集进行切分，否则将导致各进程的评估结果不一致。

1. 通常情况下，对于训练的step参数和warmup step参数，需要除以总的进程数hvd.size()，学习率需要对应增大hvd.size()倍。

说明：部分模型不需要增大学习率，如BERT模型，具体请根据训练收敛情况作判断。

step = step // hvd.size()
learning_rate = learning_rate * hvd.size()

4. 重载Optimizer。

将原单机的Optimizer使用此函数进行重载。使用示例如下：

# original optimizer
optimizer = tf.train.AdamOptimizer(learning_rate)
# AIACC-Training: wrap the Distributed Optimizer.
optimizer = hvd.DistributedOptimizer(optimizer)

DistributedOptimizer(TensorFlow_Optimizer)：用于重载Optimizer，使其可以进行AIACC-Training分布式训练。输入参数为标准TensorFlow下的Optimizer，输出为被AIACC-Training重载过的Optimizer。

5. 将当前process绑定对应的GPU卡。

• 对于TensorFlow 1.x版本：

config = tf.ConfigProto()
config.gpu_options.allow_growth= True
config.gpu_options.visible_device_list =str(hvd.local_rank())

你还需要根据以下不同情况，将config传入相应的函数中：

• 若使用tf.Session.run()，则需要在初始化tf.Session时传入config。

with tf.Session(config=config) as sess:
  sess.run(...)

• 若使用Estimator，则需要在初始化Estimator时，以tf.estimator.RunConfig(session_config=config)封装传入。

mnist_classifier = tf.estimator.Estimator(
        model_fn=cnn_model_fn, model_dir=model_dir,
        config=tf.estimator.RunConfig(session_config=config))

• 若使用Keras，需要封装后传入。

from tensorflow.keras import backend as K
K.set_session(tf.Session(config=config))

• 对于TensorFlow 2.x版本：

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')

6. 在开始训练之前，需要同步所有训练进程的初始状态，根据不同的使用形式参考以下内容进行适配。

• 对于传统的tf.Session.run()形式，需在训练前先运行tf.Session.run(hvd.broadcast_global_variables(0))。
• 对于Keras，需在训练的callbacks中增加hvd.keras.callbacks.BroadcastGlobalVariablesCallback(0)。
• 对于Estimator，需在session hook中加入hvd.BroadcastGlobalVariablesHook(0)。

7. checkpoint或模型只需要在root rank上保存，其它rank设置为none。

说明:此步骤可以防止各个进程因保存模型而发生覆盖冲突，因此务必执行。

checkpoint_dir = './checkpoints' if hvd.rank() == 0 else None

8. 通过启动perseusrun来进行训练脚本的分布式训练。以8卡2机的启动方式为例，命令示例如下：

perseusrun -np 16 -H {IP1:8},{IP2:8} python xxx.py

其中，IP为内网IP。

说明:关于启动方式的更多信息，请参见分布式启动方式说明。

示范用例

AIACC-Training软件包路径中为您提供了针对tf.Session.run()、Keras、Estimator、tf2 eager等训练方式的适配示例代码，您可以通过以下操作体验训练过程。

1. 进入适配示例代码目录。

cd `echo $(python -c "import perseus; print(perseus)") | cut -d\' -f 4 | sed "s/\_\_init\_\_\.py//"`examples/

2. 启动分布式训练。

以Keras形式的MNIST训练模型启动单机2卡为例，启动命令如下：

perseusrun -np 2 -H localhost:2 python tensorflow_keras_mnist.py

常见问题

训练过程中出现OOM（显存不足）报错

您可以识别以下几种可能情况并解决。

• 使用nvidia-smi检测启动过程中，各显卡占用显存是否均衡增加，单个process对应单个GPU，每个process所使用显存应该接近。
• 若存在多个process绑定同一张卡的情况，需要检查config.gpu_options.visible_device_list是否正确设置。
• 若存在显存直接被用尽，您可以尝试增加config.allow_growth = True。
• 若启用了XLA，您可以尝试增加参数config.gpu_options.per_process_gpu_memory_fraction = 0.9或者0.8。

如何快速判断是否是梯度通信带来的性能瓶颈

您可以将适配代码中的DistributedOptimizer(opt)参数注释掉，此时，将不会产生梯度通信，您即可进一步排查数据IO、CPU预处理等可能造成性能瓶颈的原因。

对数据集做shard的注意事项

由于AIACC-Training是由多个进程启动同一份训练代码，因此您需要对数据集做数据集切分为子数据集，使每个进程处理与训练不同的子数据集。TensorFlow为tf.data.Dataset类提供了自动切分数据的shard()接口，您可以结合size()、local_rank()或rank()函数进行自定义的数据切分，例如：

dataset = tf.dataset.shard(hvd.size(), hvd.rank())

重要:为保证算法的准确性，shard使用时必须放在数据集shuffle操作之前，且为了保证性能稳定，shard操作不要放在repeat操作之后。否则会带来额外的预处理负担，将严重影响性能。

好啦！小弹的分享到此为止。我们更欢迎您分享您对阿里云产品的设想、对功能的建议或者各种吐槽，请扫描提交问卷并获得社区积分或精美礼品一份。https://survey.aliyun.com/apps/zhiliao/P4y44bm_8

【扫码填写上方调研问卷】

欢迎每位来到弹性计算的开发者们来反馈问题哦~