具有多个GPU的分布式张量流 [英] Distributed tensorflow with multiple gpu

问题描述

tf.train.replica_device_setter 似乎不允许指定要使用的GPU。

It seems that tf.train.replica_device_setter doesn't allow specify gpu which work with.

我想做的事情如下：

 with tf.device(
   tf.train.replica_device_setter(
   worker_device='/job:worker:task:%d/gpu:%d' % (deviceindex, gpuindex)):
     <build-some-tf-graph>

推荐答案

如果您的参数是不分片，您可以使用 replica_device_setter 的简化版本，如下所示：

If your parameters are not sharded, you could do it with a simplified version of replica_device_setter like below:

def assign_to_device(worker=0, gpu=0, ps_device="/job:ps/task:0/cpu:0"):
    def _assign(op):
        node_def = op if isinstance(op, tf.NodeDef) else op.node_def
        if node_def.op == "Variable":
            return ps_device
        else:
            return "/job:worker/task:%d/gpu:%d" % (worker, gpu)
    return _assign

with tf.device(assign_to_device(1, 2)):
  # this op goes on worker 1 gpu 2
  my_op = tf.ones(())

这篇关于具有多个GPU的分布式张量流的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！