Tensorflow中ring all-reduce的实现

import tensorflow as tfimport osimport json
NUM_WORKERS = 1IP_ADDRS = ["xx.xxx.xx.xxxx", "xx.xxx.xx.xxxx"]PORTS = [2222， 2222]
def model_fn(...):    .....
def input_fn(...):    .....
# 需要每个机器配置TF_CONFIG环境变量os.environ['TF_CONFIG'] = json.dumps({    'cluster': {        'worker': ['%s:%d' % (IP_ADDRS[w], PORTS[w]) for w in range(NUM_WORKERS)]    },    'task': {'type': 'worker', 'index': 0}})
# Method for using MultiWorkerMirroredStrategystrategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
config = tf.estimator.RunConfig(train_distribute=strategy)
classifier = tf.estimator.Estimator(    model_fn=model_fn, model_dir='/tmp/multiworker', config=config)tf.estimator.train_and_evaluate(    classifier,    train_spec=tf.estimator.TrainSpec(input_fn=input_fn),    eval_spec=tf.estimator.EvalSpec(input_fn=input_fn))

import osimport errnoimport tensorflow as tfimport horovod.tensorflow as hvdimport numpy as npfrom tensorflow import keras
tf.logging.set_verbosity(tf.logging.INFO)
def create_model(feature, target, mode):    .....
def train_input_generator(x_train, y_train, batch_size=64):    .....
def main(_):    # Horovod: initialize Horovod.    hvd.init()
    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' % hvd.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)
    # Horovod: adjust learning rate based on number of GPUs.    opt = tf.train.RMSPropOptimizer(0.001 * hvd.size())
    # Horovod: add Horovod Distributed Optimizer.    opt = hvd.DistributedOptimizer(opt)
    global_step = tf.train.get_or_create_global_step()    train_op = opt.minimize(loss, global_step=global_step)
    hooks = [        # Horovod: 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.        hvd.BroadcastGlobalVariablesHook(0),
        # Horovod: adjust number of steps based on number of GPUs.        tf.train.StopAtStepHook(last_step=20000 // hvd.size()),
        tf.train.LoggingTensorHook(tensors={'step': global_step, 'loss': loss},                                   every_n_iter=10),    ]
    # Horovod: 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(hvd.local_rank())
    # Horovod: save checkpoints only on worker 0 to prevent other workers from    # corrupting them.    checkpoint_dir = './checkpoints' if hvd.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_})

if __name__ == "__main__":    tf.app.run()2. horovod搭配estimator
from __future__ import absolute_importfrom __future__ import divisionfrom __future__ import print_function
import osimport errnoimport numpy as npimport tensorflow as tfimport horovod.tensorflow as hvd
from tensorflow import keras
tf.logging.set_verbosity(tf.logging.INFO)
def create_model(input_shape, num_classes):     .....


def cnn_model_fn(features, labels, mode):    model = create_model()    logits = model(features)    predictions = {        # Generate predictions (for PREDICT and EVAL mode)        "classes": tf.argmax(input=logits, axis=1),        # Add `softmax_tensor` to the graph. It is used for PREDICT and by the        # `logging_hook`.        "probabilities": tf.nn.softmax(logits, name="softmax_tensor")    }    if mode == tf.estimator.ModeKeys.PREDICT:        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
    # Calculate Loss (for both TRAIN and EVAL modes)    onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10)    loss = tf.losses.softmax_cross_entropy(        onehot_labels=onehot_labels, logits=logits)
    # Configure the Training Op (for TRAIN mode)    if mode == tf.estimator.ModeKeys.TRAIN:        # Horovod: scale learning rate by the number of workers.        optimizer = tf.train.MomentumOptimizer(            learning_rate=0.001 * hvd.size(), momentum=0.9)
        # Horovod: add Horovod Distributed Optimizer.        optimizer = hvd.DistributedOptimizer(optimizer)
        train_op = optimizer.minimize(            loss=loss,            global_step=tf.train.get_global_step())        return tf.estimator.EstimatorSpec(mode=mode, loss=loss,                                          train_op=train_op)
    # Add evaluation metrics (for EVAL mode)    eval_metric_ops = {        "accuracy": tf.metrics.accuracy(            labels=labels, predictions=predictions["classes"])}    return tf.estimator.EstimatorSpec(        mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)

def main(unused_argv):    # Horovod: initialize Horovod.    hvd.init()
    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.    (train_data, train_labels), (eval_data, eval_labels) = \        keras.datasets.mnist.load_data('MNIST-data-%d' % hvd.rank())
    train_data = np.reshape(train_data, (-1, 784)) / 255.0    eval_data = np.reshape(eval_data, (-1, 784)) / 255.0
    # Horovod: 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(hvd.local_rank())
    # Horovod: save checkpoints only on worker 0 to prevent other workers from    # corrupting them.    model_dir = './mnist_convnet_model' if hvd.rank() == 0 else None
    # Create the Estimator    mnist_classifier = tf.estimator.Estimator(        model_fn=cnn_model_fn, model_dir=model_dir,        config=tf.estimator.RunConfig(session_config=config))
    # Set up logging for predictions    # Log the values in the "Softmax" tensor with label "probabilities"    tensors_to_log = {"probabilities": "softmax_tensor"}    logging_hook = tf.train.LoggingTensorHook(        tensors=tensors_to_log, every_n_iter=500)
    # Horovod: 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.    bcast_hook = hvd.BroadcastGlobalVariablesHook(0)
    # Train the model    train_input_fn = tf.estimator.inputs.numpy_input_fn(        x={"x": train_data},        y=train_labels,        batch_size=100,        num_epochs=None,        shuffle=True)
    # Horovod: adjust number of steps based on number of GPUs.    mnist_classifier.train(        input_fn=train_input_fn,        steps=20000 // hvd.size(),        hooks=[logging_hook, bcast_hook])
    # Evaluate the model and print results    eval_input_fn = tf.estimator.inputs.numpy_input_fn(        x={"x": eval_data},        y=eval_labels,        num_epochs=1,        shuffle=False)    eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn)    print(eval_results)

if __name__ == "__main__":    tf.app.run()

horovodrun -np 16 -H server1:4,server2:4,server3:4,server4:4 python train.py