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

Question

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

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

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

Answer 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 squeezenet模型在具有4个GPU的机器上工作.我的修改在这里.