tensorflow使用horovod進行多gpu訓練

tensorflow使用horovod多gpu訓練

要使用Horovod，在程序中添加以下內容。此示例使用TensorFlow。

1. 運行hvd.init()
   

 

2. 使用固定服務器GPU，以供此過程使用config.gpu_options.visible_device_list
   

   通過每個進程一個GPU的典型設置，您可以將其設置為local rank。在這種情況下，服務器上的第一個進程將被分配第一GPU，第二個進程將被分配第二GPU，依此類推。

 

3. 通過工人人數來衡量學習率

   同步分布式培訓中的有效批處理規模是根據工人人數來衡量的。學習率的提高彌補了批量大小的增加。

 

4. 將優化器包裝在中hvd.DistributedOptimizer
   

   分布式優化器將梯度計算委派給原始優化器，使用allreduce或allgather對梯度求平均，然后應用這些平均梯度。

 

5. 添加hvd.BroadcastGlobalVariablesHook(0)到播放初始變量狀態從0級到所有其他進程

   當使用隨機權重開始訓練或從檢查點恢復訓練時，這是確保所有工人進行一致初始化的必要步驟。另外，如果您不使用MonitoredTrainingSession，則可以hvd.broadcast_global_variables在初始化全局變量之后執行op。

 

6. 修改您的代碼以僅在工作程序0上保存檢查點，以防止其他工作程序破壞它們

   通過傳遞checkpoint_dir=None給tf.train.MonitoredTrainingSessionif 來完成此操作。hvd.rank() != 0

簡單示例代碼

import tensorflow as tf
import horovod.tensorflow as hvd


# Initialize Horovod
hvd.init()

# Pin GPU to be used to process local rank (one GPU per process)
config = tf.ConfigProto()
config.gpu_options.visible_device_list = str(hvd.local_rank())

# Build model...
loss = ...
opt = tf.train.AdagradOptimizer(0.01 * hvd.size())

# Add Horovod Distributed Optimizer
opt = hvd.DistributedOptimizer(opt)

# Add hook to broadcast variables from rank 0 to all other processes during
# initialization.
hooks = [hvd.BroadcastGlobalVariablesHook(0)]

# Make training operation
train_op = opt.minimize(loss)

# Save checkpoints only on worker 0 to prevent other workers from corrupting them.
checkpoint_dir = '/tmp/train_logs' if hvd.rank() == 0 else None

# The MonitoredTrainingSession takes care of session initialization,
# restoring from a checkpoint, saving to a checkpoint, and closing when done
# or an error occurs.
with tf.train.MonitoredTrainingSession(checkpoint_dir=checkpoint_dir,
                                       config=config,
                                       hooks=hooks) as mon_sess:
  while not mon_sess.should_stop():
    # Perform synchronous training.
    mon_sess.run(train_op)