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

时间:2017-11-10 13:22:04

标签: tensorflow tensorflow-gpu multi-gpu

我有一个标准的tensorflow Estimator和一些模型,并希望在多个GPU而不是一个GPU上运行它。如何使用数据并行来完成?

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

是否有人使用tf.learn.Estimator有一个很好的例子?或指向教程的链接?

5 个答案:

答案 0 :(得分:5)

我认为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.TowerOptimizemodel_fn()将优化程序包装为tf.contrib.estimator.replicate_model_fn()。 我按照说明进行操作,使TPU squeezenet模型在具有4个GPU的机器上工作。我的修改here

答案 1 :(得分:1)

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

运行数据并行的一种方法是循环遍历可用的GPU设备,并将批量块发送到模型的复制版本(全部在model_fn中完成),然后合并结果。

答案 2 :(得分:1)

您可以使用范围和设备:

 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

答案 3 :(得分:0)

我认为这就是您所需要的。

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

更多详细信息:https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy

解释: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)

答案 4 :(得分:0)

您可以找到使用var value; if(applications.gecko.hasOwnProperty('id')) value = applications.gecko.id; else value = 'error'; tf.distribute.MirroredStrategy here的示例。