TensorFlow深度学习并发加速训练-CFANZ编程社区

TensorFlow深度学习并发训练
TensorFlow 2.0 分布式训练

单机多卡训练: MirroredStrategy

tf.distribute.MirroredStrategy 是一种简洁且高性能的，数据并行的同步式分布式策略，主要支持多个 GPU 在同一台主机上训练。使用这种策略时，我们只需实例化一个 MirroredStrategy 策略:
strategy = tf.distribute.MirroredStrategy()
并将模型构建的代码放入 strategy.scope() 的上下文环境中:
with strategy.scope():
    # 模型构建代码
小技巧
可以在参数中指定设备，如:
strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
即指定只使用第 0、1

我们使用同一台主机上的 4 块 NVIDIA GeForce GTX 1080 Ti 显卡进行单机多卡的模型训练。所有测试的 epoch 数均为 5。使用单机无分布式配置时，虽然机器依然具有 4 块显卡，但程序不使用分布式的设置，直接进行训练，Batch Size 设置为 64。使用单机四卡时，测试总 Batch Size 为 64（分发到单台机器的 Batch Size 为 16）和总 Batch Size 为 256（分发到单台机器的 Batch Size 为 64）两种情况

TensorFlow深度学习并发加速训练_3d

可见，使用 MirroredStrategy 后，模型训练的速度有了大幅度的提高。在所有显卡性能接近的情况下，训练时长与显卡的数目接近于反比关系。

MirroredStrategy 过程简介
MirroredStrategy 的步骤如下：

训练开始前，该策略在所有 N 个计算设备上均各复制一份完整的模型；
每次训练传入一个批次的数据时，将数据分成 N 份，分别传入 N 个计算设备（即数据并行）；

N 个计算设备使用本地变量（镜像变量）分别计算自己所获得的部分数据的梯度；

使用分布式计算的 All-reduce 操作，在计算设备间高效交换梯度数据并进行求和，使得最终每个设备都有了所有设备的梯度之和；

使用梯度求和的结果更新本地变量（镜像变量）；

当所有设备均更新本地变量后，进行下一轮训练（即该并行策略是同步的）。

默认情况下，TensorFlow 中的 MirroredStrategy 策略使用 NVIDIA NCCL 进行 All-reduce

多机训练:MultiWorkerMirroredStrategy

多机训练的方法和单机多卡类似，将 MirroredStrategy 更换为适合多机训练的 MultiWorkerMirroredStrategy 即可。不过，由于涉及到多台计算机之间的通讯，还需要进行一些额外的设置。具体而言，需要设置环境变量 TF_CONFIG ，示例如下:

os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["localhost:20000", "localhost:20001"]
    },
    'task': {'type': 'worker', 'index': 0}
})
TF_CONFIG 由 cluster 和 task 两部分组成：
cluster 说明了整个多机集群的结构和每台机器的网络地址（IP + 端口号）。对于每一台机器，cluster 的值都是相同的；
task 说明了当前机器的角色。例如，{'type': 'worker', 'index': 0} 说明当前机器是 cluster 中的第 0 个 worker（即 localhost:20000

分别建立两台具有单张 NVIDIA Tesla K80 的虚拟机实例（具体建立方式参见后文介绍），并分别测试在使用一个 GPU 时的训练时长和使用两台虚拟机实例进行分布式训练的训练时长。所有测试的 epoch 数均为 5。使用单机单卡时，Batch Size 设置为 64。使用双机单卡时，测试总 Batch Size 为 64（分发到单台机器的 Batch Size 为 32）和总 Batch Size 为 128（分发到单台机器的 Batch Size 为 64）两种情况。

TensorFlow深度学习并发加速训练_3d_02

可见模型训练的速度同样有大幅度的提高。在所有机器性能接近的情况下，训练时长与机器的数目接近于反比关系。

参考：

https://mp.weixin.qq.com/s?subscene=23&__biz=MzU1OTMyNDcxMQ==&mid=2247487445&idx=2&sn=e0aa4dfc07d83a1fce08ca07cd0f754e&chksm=fc18449dcb6fcd8b0105bef0af4e26aec7b212deeb487818e436fc2df8efa5acc08e13f70879&scene=7&key=008658d21f0cfb81470be9aca4f1876a68d8b5db0af3a017cb41d0490263d26d18be273445110052dcfbe6e31b7335d9cbf3ed4b022d4b299958c2e6a323ff340e5db09e1711db1b55b0a60c2ecd3af7&ascene=0&uin=MjMzMTY4OTI2MQ%3D%3D&devicetype=Windows+10&versinotallow=6208006f&lang=zh_CN&exportkey=AQsP3mGaPPoKT2Pmc17vIgA%3D&pass_ticket=a4lCBHXReXlR9%2B64n8lmukNzyq5Bv62HAf42XXwkUpT4Okokk7pURRgtCi0c4Qie

切分GPU资源

管理 GPU 内存
默认情况下，TensorFlow 会在您第一次运行图形时自动获取所有可用 GPU 中的所有 RAM，因此当第一个程序仍在运行时，您将无法启动第二个 TensorFlow 程序。如果你尝试，会报错：

****一种解决方案是在不同的 GPU 卡上运行每个进程。为此，最简单的选择是设置CUDA_VISIBLE_DEVICES环境变量，以便每个进程只能看到对应的 GPU 卡。例如，你可以像这样启动两个程序：

TensorFlow深度学习并发加速训练_3d_03

另一种选择是告诉 TensorFlow 只抓取一小部分内存。例如，要使 TensorFlow 只占用每个 GPU 内存的 40%，您必须创建一个ConfigProto对象，将其gpu_options.per_process_gpu_memory_fraction选项设置为 0.4，并使用以下配置创建session：

TensorFlow深度学习并发加速训练_数据_04

TensorFlow深度学习并发加速训练_tensorflow_05

另一种选择是告诉 TensorFlow 只在需要时才抓取内存。为此，您必须将config.gpu_options.allow_growth设置为True。但是，TensorFlow 一旦抓取内存就不会释放内存（以避免内存碎片），因此您可能会在一段时间后内存不足。是否使用此选项可能难以确定，因此一般而言，您可能想要坚持之前的某个选项。

参考：

https://yq.aliyun.com/articles/603783?utm_cnotallow=m_1000003941