分布式Tensorflow设备在Google引擎中的放置
分布式Tensorflow设备在Google引擎中的放置
提问于 2017-12-13 10:59:23
我在google云ML引擎中运行一个大型分布式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)说
回答 1
Stack Overflow用户
回答已采纳
发布于 2017-12-13 12:08:55
除了变量之外，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页面原文内容由Stack Overflow提供。腾讯云小微IT领域专用引擎提供翻译支持
原文链接：
https://stackoverflow.com/questions/47791372
复制相关文章
相似问题
腾讯云开发者
