TensorFlow 2.0教程05：跨多个节点的分布式培训

分布式训练允许扩大深度学习任务，因此可以学习更大的models或以更快的速度进行训练。在之前的教程中，我们讨论了如何MirroredStrategy在单个节点（物理机器）内实现多GPU训练。在本教程中，我们将解释如何在多个节点之间进行分布式训练。本教程包括：

• 用于多节点分布式培训的代码样板。
• 示例代码在多台计算机上运行。

要重现本教程，请参考TensorFlow 2分布式培训 github repository。

 

代码样板

与单节点内的多GPU训练相似，多节点训练也使用分布式策略。在这种情况下，tf.distribute.experimental.MultiWorkerMirroredStrategy。多节点训练还需要设置TF_CONFIG环境变量。请注意，每个节点上的环境变量都会略有不同。例如，这是worker 0两节点分布式培训作业中的设置：

 

os.environ["TF_CONFIG"] = json.dumps({
    'cluster': {
        'worker': ["10.1.10.58:12345", "10.1.10.250:12345"]
    },
    'task': {'type': 'worker', 'index': 0}
})

 

本质上，TF_CONFIG是一个JSON字符串，代表集群并标识该机器在该集群中的角色。上面的代码设置了TF_CONFIG环境变量，也可以使用命令行导出或将其作为shell命令的前缀进行设置，例如：

TF_CONFIG='{"cluster": {"worker": ["10.1.10.58:12345", "10.1.10.250:12345"]}, "task": {"index": 0, "type": "worker"}}' python worker.py

 

cluster所有节点上的字段均相同。它描述了如何设置群集。在这种情况下，我们的集群只有两个工作节点，其IP:port信息列在worker阵列中。该task场从节点到不同的节点。它指定节点的类型和索引，然后将其用于从cluster字段中获取详细信息，并提供给任务管理器以划分工作。在这种情况下，此配置文件指示培训作业在工作程序0上运行，即"10.1.10.58:12345"

我们需要为每个节点自定义此python代码段。因此，第二个节点将具有'task': {'type': 'worker', 'index': 1}。

然后，我们需要创建分布式策略：

strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()

 

注意，此行必须TF_CONFIG在定义数据管道和模型之后以及在定义数据管道和模型之前完成。否则，Collective ops must be configured at program startup将触发错误。

代码样板的最后一部分定义了策略范围内的模型：

with strategy.scope():
  model = resnet.resnet56(img_input=img_input, classes=NUM_CLASSES)
  model.compile(
            optimizer=opt,
            loss='sparse_categorical_crossentropy',
            metrics=['sparse_categorical_accuracy']) 
model.fit(train_dataset,
          epochs=NUM_EPOCHS)

 

 

进行训练

要运行分布式培训，需要自定义培训脚本并将其复制到所有节点。为了更加清楚，我们可以使用前缀语法设置环境变量。各个节点的设置不同。

确保节点可以在不要求密码的情况下相互连接。最方便的方法是使用ssh密钥而不是密码身份验证。如何使用ssh键。

最后，在两个节点上同时运行脚本。

# On the first nodeTF_CONFIG='{"cluster": {"worker": ["10.1.10.58:12345", "10.1.10.250:12345"]}, "task": {"index": 0, "type": "worker"}}' python worker.py# On the second nodeTF_CONFIG='{"cluster": {"worker": ["10.1.10.58:12345", "10.1.10.250:12345"]}, "task": {"index": 1, "type": "worker"}}' python worker.py

训练现在分布在多个节点上。由于使用了该Mirrored策略，因此两个节点的输出已同步。

 

总结

本教程说明了如何在TensorFlow 2中进行分布式训练。关键是设置TF_CONFIG环境变量并使用MultiWorkerMirroredStrategy来限定模型定义的范围。

在本教程中，我们需要在每个节点上手动运行定制的TF_CONFIG训练脚本。可以看到，随着节点数量的增加，设置环境变量很快变得乏味。有更高级的方法可以在大量节点上部署分布式培训作业，例如Horovod，带有TF-flow的Kubernetes，OpenMPI或使用诸如Ansible的部署脚本。要在本教程中重现结果，请参考此TensorFlow 2分布式培训教程github repository.

接下来，给大家介绍一下租用GPU做实验的方法，我们是在智星云租用的GPU，使用体验很好。具体大家可以参考：智星云官网： http://www.ai-galaxy.cn/，淘宝店：https://shop36573300.taobao.com/公众号: 智星AI

 

参考文献：

https://lambdalabs.com/blog/tensorflow-2-0-tutorial-05-distributed-training-multi-node/

https://github.com/lambdal/TensorFlow2-tutorial/tree/master/05-distributed-training

https://debian-administration.org/article/530/SSH_with_authentication_key_instead_of_password​