horovod + tf.train.CheckpointSaverHook

最近在做分布式模型训练相关工作,利用到了horovod框架,当horovod+tf.train.MonitoredTrainingSession+tf.train.CheckpointSaverHook时,会出现horovod rank抢占之类的报错。并且在log中多次出现Create CheckpointSaverHook的信息。

并且由于MonitoredTrainingSession的重启机制,session重复start,报错也一直重复。

相关代码如下:

...
checkpoint_path = ... if  hvd.rank() == 0 else None
hook_worker = [hvd.BroadcastGlobalVariablesHook(0)]
hook_master = [
               tf.train.CheckpointSaverHook(checkpoint_dir,
                                            save_steps,
                                            saver,
                                            checkpoint_basename='model.ckpt')
              ]
hooks = hook_worker + hook_master if hvd.rank()==0 else hook_worker
with tf.train.MonitoredTrainingSession(checkpoint_dir = None,
                                       config = config,
                                       hooks = hooks
                                       ) as mon_sess:
...

之后将这一部分代码改为:

hook_worker = [
               hvd.BroadcastGlobalVariablesHook(0)
              ]

if hvd.rank()==0:
  hook_master = [
                 tf.train.CheckpointSaverHook(...),
                 tf.train.SummarySaverHook(...)
                ]
else:
  hook_master = []

hooks = hook_worker + hook_master

这一部分的错误可能和CheckpointSaverHook源码有关,需要进一步研究源码。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值