tensorflow 保存训练loss_【分布式训练】单机多卡的正确打开方式（四）：Horovod...

最新推荐文章于 2022-05-10 22:56:11 发布

臭熊的哥哥

最新推荐文章于 2022-05-10 22:56:11 发布

阅读量562

点赞数

文章标签： tensorflow 保存训练loss

本文链接：https://blog.csdn.net/weixin_42519488/article/details/113628339

版权

本文介绍了Uber开源的分布式训练工具Horovod，适用于TensorFlow、Keras和PyTorch。Horovod采用Ring-AllReduce算法，简化了多GPU训练，并提供了TensorFlow和PyTorch的示例。文章强调了有效批次大小、学习率调整、数据分片等关键点，并提及Horovod与直接使用框架API的比较。

摘要由CSDN通过智能技术生成

讲完了单机多卡的分布式训练的理论、TensorFlow和PyTorch分别的实现后，今天瓦砾讲一个强大的第三方插件：Horovod。

Horovod是Uber开源的跨平台的分布式训练工具，名字来自于俄国传统民间舞蹈，舞者手牵手围成一个圈跳舞，与Horovod设备之间的通信模式很像，有以下几个特点：

兼容TensorFlow、Keras和PyTorch机器学习框架。
使用Ring-AllReduce算法，对比Parameter Server算法，有着无需等待，负载均衡的优点。
实现简单，五分钟包教包会。（划重点）

Uber官方在git上给了很详细的例子： https://github.com/horovod/horovod/tree/master/examples，所以这里只简单讲一下大概的使用方法：

TensorFlow

以TF的Custom Training Loop API为例：

import tensorflow as tf
import horovod.tensorflow as hvd

# 1. 初始化horovod
hvd.init()
# 2. 给当前进程分配对应的gpu，local_rank()返回的是当前是第几个进程
config = tf.ConfigProto()
config.gpu_options.visible_device_list = str(hvd.local_rank())
# 3. Scale学习率，封装优化器
opt = tf.train.AdagradOptimizer(0.01 * hvd.size())
opt = hvd.DistributedOptimizer(opt)
# 4. 定义初始化的时候广播参数的hook，这个是为了在一开始的时候同步各个gpu之间的参数
hooks = [hvd.BroadcastGlobalVariablesHook(0)]
# 搭建model，定义loss
loss = ...
train_op = opt.minimize(loss)
# 5. 只保存一份ckpt就行
checkpoint_dir = '/tmp/train_logs' if hvd.rank() == 0 else None
# 7. 用MonitoredTrainingSession实现初始化，读写ckpt
with tf.train.MonitoredTrainingSession(checkpoint_dir=checkpoint_dir,
                                       co