tensorflow - 如何在具有数据并行性的多个 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() .
我遵循了描述，并在具有 4 个 GPU 的机器上制作了一个 TPU 挤压网模型。我的修改 here .