code: def main(_):
      # BytePS: initialize BytePS.
      bps.init()
  
      # Keras automatically creates a cache directory in ~/.keras/datasets for
      # storing the downloaded MNIST data. This creates a race
      # condition among the workers that share the same filesystem. If the
      # directory already exists by the time this worker gets around to creating
      # it, ignore the resulting exception and continue.
      cache_dir = os.path.join(
          os.path.expanduser("~"), ".keras", "datasets"
      )
      if not os.path.exists(cache_dir):
          try:
              os.mkdir(cache_dir)
          except OSError as e:
              if e.errno == errno.EEXIST and os.path.isdir(cache_dir):
                  pass
              else:
                  raise
  
      # Download and load MNIST dataset.
      (x_train, y_train), (
          x_test,
          y_test,
      ) = keras.datasets.mnist.load_data("MNIST-data-%d" % bps.rank())
  
      # The shape of downloaded data is (-1, 28, 28), hence we need to reshape it
      # into (-1, 784) to feed into our network. Also, need to normalize the
      # features between 0 and 1.
      x_train = np.reshape(x_train, (-1, 784)) / 255.0
      x_test = np.reshape(x_test, (-1, 784)) / 255.0
  
      # Build model...
      with tf.name_scope("input"):
          image = tf.placeholder(tf.float32, [None, 784], name="image")
          label = tf.placeholder(tf.float32, [None], name="label")
      predict, loss = conv_model(
          image, label, tf.estimator.ModeKeys.TRAIN
      )
  
      # BytePS: adjust learning rate based on number of GPUs.
      opt = tf.train.RMSPropOptimizer(0.001 * bps.size())
  
      # BytePS: add BytePS Distributed Optimizer.
      opt = bps.DistributedOptimizer(opt)
  
      global_step = tf.train.get_or_create_global_step()
      train_op = opt.minimize(loss, global_step=global_step)
  
      hooks = [
          # BytePS: BroadcastGlobalVariablesHook broadcasts initial variable states
          # from rank 0 to all other processes. This is necessary to ensure consistent
          # initialization of all workers when training is started with random weights
          # or restored from a checkpoint.
          bps.BroadcastGlobalVariablesHook(0),
          # BytePS: adjust number of steps based on number of GPUs.
          tf.train.StopAtStepHook(last_step=200000 // bps.size()),
          tf.train.LoggingTensorHook(
              tensors={"step": global_step, "loss": loss},
              every_n_iter=10,
          ),
      ]
  
      # BytePS: pin GPU to be used to process local rank (one GPU per process)
      config = tf.ConfigProto()
      config.gpu_options.allow_growth = True
      config.gpu_options.visible_device_list = str(bps.local_rank())
  
      # BytePS: save checkpoints only on worker 0 to prevent other workers from
      # corrupting them.
      checkpoint_dir = "./checkpoints" if bps.rank() == 0 else None
      training_batch_generator = train_input_generator(
          x_train, y_train, batch_size=100
      )
      # 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, hooks=hooks, config=config
      ) as mon_sess:
          while not mon_sess.should_stop():
              # Run a training step synchronously.
              image_, label_ = next(training_batch_generator)
              mon_sess.run(
                  train_op, feed_dict={image: image_, label: label_}
              )
  
file_path: byteps/example/tensorflow/tensorflow_mnist.py
imports: import os
  import errno
  import byteps.tensorflow
  import numpy
  from tensorflow import keras
  import tensorflow
label_desc: Use image node_id 3 for calling a global function with example usage: main(_) without return types