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

Question

我是运行 google 云 ML 引擎中的大型分布式 Tensorflow 模型。我想使用带 GPU 的机器。我的图表由两个主要部分组成，即 input/data reader 函数和计算部分。

我希望在 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

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

Distributed Tensorflow device placement in Google Cloud ML engine

gpu

distributed-computing

multi-gpu

tensorflow

google-cloud-ml