Distributed tensorflow with multiple gpu

Question

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

What I want to do is like below:

 with tf.devi         


        

Answer 1

I didn't check previous versions, but in Tensorflow 1.4/1.5, you can specify devices in replica_device_setter(worker_device='job:worker/task:%d/gpu:%d' % (FLAGS.task_index, i), cluster=self.cluster).

See tensorflow/python/training/device_setter.py line 199-202:

if ps_ops is None: # TODO(sherrym): Variables in the LOCAL_VARIABLES collection should not be # placed in the parameter server. ps_ops = ["Variable", "VariableV2", "VarHandleOp"]

Thanks to the code provided by @Yaroslav Bulatov, but his protocol is different from replica_device_setter, and may fail in some cases.