1 import tensorflow as tf
 2 import horovod.tensorflow as hvd
 3 
 4 
 5 # Initialize Horovod
 6 hvd.init()
 7 
 8 # Pin GPU to be used to process local rank (one GPU per process)
 9 config = tf.ConfigProto()
10 config.gpu_options.visible_device_list = str(hvd.local_rank())
11 
12 # Build model...
13 loss = ...
14 opt = tf.train.AdagradOptimizer(0.01 * hvd.size())
15 
16 # Add Horovod Distributed Optimizer
17 opt = hvd.DistributedOptimizer(opt)
18 
19 # Add hook to broadcast variables from rank 0 to all other processes during
20 # initialization.
21 hooks = [hvd.BroadcastGlobalVariablesHook(0)]
22 
23 # Make training operation
24 train_op = opt.minimize(loss)
25 
26 # Save checkpoints only on worker 0 to prevent other workers from corrupting them.
27 checkpoint_dir = ‘/tmp/train_logs‘ if hvd.rank() == 0 else None
28 
29 # The MonitoredTrainingSession takes care of session initialization,
30 # restoring from a checkpoint, saving to a checkpoint, and closing when done
31 # or an error occurs.
32 with tf.train.MonitoredTrainingSession(checkpoint_dir=checkpoint_dir,
33                                        config=config,
34                                        hooks=hooks) as mon_sess:
35   while not mon_sess.should_stop():
36     # Perform synchronous training.
37     mon_sess.run(train_op)