tensorflow - 如何在 tf contrib estimator 中使用 GPU

Question

我在云 ml 上有一台带 4 个 GPU 的机器。如何让我的 tf contrib 估算器使用机器上的 GPU？我可以让估算器使用所有这些吗？

Answer 1

有两种类型的估算器:“pre-made ”和“custom ”。

没有任何预制估算器(例如 DNNClassifier)，目前可以与多个 GPU 一起工作——所有计算都将在单个 GPU 上进行。

UPDATE 2018/07/18*(基于墓碑的帖子):DistributionStrategy简化了创建在多个 GPU 上运行的“自定义”估算器的过程。

“自定义”估算器是您编写的估算器。因此，他们可以做任何您想做的事情，包括利用多个 GPU。在这种情况下，您只需确保使用适当的 tf.device 语句编写您的 model_fn 以将操作分配给正确的 GPU(TensorFlow requires 您手动当你有多个 GPU 时分配操作)。这通常涉及使用 towers ，这只是意味着在每个 GPU 上重复模型结构。来自 TensorFlow docs :

# Creates a graph.
c = []
for d in ['/gpu:2', '/gpu:3']:
  with tf.device(d):
    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3])
    b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2])
    c.append(tf.matmul(a, b))

以上通常效果很好，但如果您想进一步优化性能，请参阅此 advanced guide .

CIFAR-10 sample演示了在估计器中使用多个 GPU 进行图像分类。