tensorflow - Google Cloud ML 引擎中的分布式 Tensorflow 设备放置

Question

我在谷歌云机器学习引擎中运行一个大型分布式 Tensorflow 模型。我想使用带有 GPU 的机器。
我的图形由两个主要部分组成，输入/数据读取器功能和计算部分。

我希望在 PS 任务中放置变量，在 CPU 中的输入部分和在 GPU 上的计算部分。
函数tf.train.replica_device_setter自动将变量放置在 PS 服务器中。

这是我的代码的样子:

with tf.device(tf.train.replica_device_setter(cluster=cluster_spec)):
    input_tensors = model.input_fn(...)
    output_tensors = model.model_fn(input_tensors, ...)

是否可以使用 tf.device()连同replica_device_setter()如:

with tf.device(tf.train.replica_device_setter(cluster=cluster_spec)):
    with tf.device('/cpu:0')
        input_tensors = model.input_fn(...)
    with tf.device('/gpu:0')
        tensor_dict = model.model_fn(input_tensors, ...)

请问replica_divice_setter()被覆盖并且变量未放置在 PS 服务器中？

此外，由于集群中的设备名称类似于 job:master/replica:0/task:0/gpu:0我怎么说 Tensorflow tf.device(whatever/gpu:0) ?

Answer 1

在 tf.train.replica_device_setter 中的任何操作，变量之外块自动固定到 "/job:worker" ，这将默认为“worker”作业中第一个任务管理的第一个设备。

您可以 使用嵌入式设备块将它们固定到另一个设备(或任务):

with tf.device(tf.train.replica_device_setter(ps_tasks=2, ps_device="/job:ps", 
                                          worker_device="/job:worker")):
  v1 = tf.Variable(1., name="v1")  # pinned to /job:ps/task:0 (defaults to /cpu:0)
  v2 = tf.Variable(2., name="v2")  # pinned to /job:ps/task:1 (defaults to /cpu:0)
  v3 = tf.Variable(3., name="v3")  # pinned to /job:ps/task:0 (defaults to /cpu:0)
  s = v1 + v2            # pinned to /job:worker (defaults to task:0/cpu:0)
  with tf.device("/task:1"):
    p1 = 2 * s           # pinned to /job:worker/task:1 (defaults to /cpu:0)
    with tf.device("/cpu:0"):
      p2 = 3 * s         # pinned to /job:worker/task:1/cpu:0