分布式训练允许扩大深度学习任务,因此可以学习更大的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