如何在具有数据并行性的多个 GPU 上运行 Tensorflow Estimator答案

【问题标题】：How to run Tensorflow Estimator on multiple GPUs with data parallelism如何在具有数据并行性的多个 GPU 上运行 Tensorflow Estimator
【发布时间】：2017-11-10 13:22:03
【问题描述】：

我有一个带有一些模型的标准 tensorflow Estimator，并希望在多个 GPU 上运行它，而不仅仅是一个。如何使用数据并行性来做到这一点？

我搜索了 Tensorflow 文档，但没有找到示例；只有句子说使用 Estimator 会很容易。

有人有使用 tf.learn.Estimator 的好例子吗？或者教程的链接？

【问题讨论】：

我认为tf.contrib.estimator.replicate_model_fn 是更清洁的解决方案。我按照描述使 TPU 挤压网络模型在具有 4 个 GPU 的机器上工作。我的修改here.

标签： tensorflow tensorflow-gpu multi-gpu

【解决方案1】：

我认为tf.contrib.estimator.replicate_model_fn 是一个更清洁的解决方案。以下来自tf.contrib.estimator.replicate_model_fn文档，

...
def model_fn(...):  # See `model_fn` in `Estimator`.
  loss = ...
  optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
  optimizer = tf.contrib.estimator.TowerOptimizer(optimizer)
  if mode == tf.estimator.ModeKeys.TRAIN:
    #  See the section below on `EstimatorSpec.train_op`.
    return EstimatorSpec(mode=mode, loss=loss,
                         train_op=optimizer.minimize(loss))

  #  No change for `ModeKeys.EVAL` or `ModeKeys.PREDICT`.
  return EstimatorSpec(...)
...
classifier = tf.estimator.Estimator(
  model_fn=tf.contrib.estimator.replicate_model_fn(model_fn))

您需要做的是用tf.contrib.estimator.TowerOptimize 和model_fn() 用tf.contrib.estimator.replicate_model_fn() 包装优化器。我按照描述使 TPU 挤压网络模型在具有 4 个 GPU 的机器上工作。我的修改here.

【讨论】：

它似乎被删除了。 THIS FUNCTION IS DEPRECATED. It will be removed after 2018-05-31. Instructions for updating: Please use tf.contrib.distribute.MirroredStrategy instead.tensorflow.org/api_docs/python/tf/contrib/estimator/…

【解决方案2】：

我想这就是你所需要的。

链接：https://www.youtube.com/watch?v=bRMGoPqsn20

解释：https://medium.com/tensorflow/multi-gpu-training-with-estimators-tf-keras-and-tf-data-ba584c3134db

NUM_GPUS = 8
dist_strategy = tf.contrib.distribute.MirroredStrategy(num_gpus=NUM_GPUS)
config = tf.estimator.RunConfig(train_distribute=dist_strategy)
estimator = tf.estimator.Estimator(model_fn,model_dir,config=config)

更新

对于 TF-2.0 和 Keras，您可以使用它 (https://www.tensorflow.org/tutorials/distribute/keras)

【讨论】：

【解决方案3】：

标准例子是：https://github.com/tensorflow/tensorflow/blob/r1.4/tensorflow/contrib/learn/python/learn/estimators/estimator.py

以数据并行方式运行它的一种方法是遍历可用的 GPU 设备，并将批量发送到模型的复制版本（全部在您的 model_fn 中完成），然后合并结果。

【讨论】：

【解决方案4】：

您可以为此使用范围和设备：

 with tf.variable_scope(tf.get_variable_scope()):
  for i in xrange(FLAGS.num_gpus):
    with tf.device('/gpu:%d' % i):
      with tf.name_scope('%s_%d' % (cifar10.TOWER_NAME, i)) as scope:

那里有完整的例子： https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_multi_gpu_train.py

【讨论】：

【解决方案5】：

您可以找到使用tf.distribute.MirroredStrategy 和tf.estimator.train_and_evaluate here 的示例。

【讨论】：