在tensorflow中，colocate_用于什么？

Question

这是官方文件的链接。 https://www.tensorflow.org/versions/r1.3/api_docs/python/tf/colocate_with

Answer 1

它是一个上下文管理器，用于确保将要创建的操作或张量放置在引用操作所在的同一设备上。考虑这段代码（经过测试）：

import tensorflow as tf

with tf.device("/cpu:0"):
  a = tf.constant(0.0, name="a")

with tf.device("/gpu:0"):
  b = tf.constant(0.0, name="b")
  with tf.colocate_with(a):
    c = tf.constant(0.0, name="c")
  d = tf.constant(0.0, name="d")

for operation in tf.get_default_graph().get_operations():
  print(operation.name, operation.device)

输出：

(u'a', u'/device:CPU:0')
(u'b', u'/device:GPU:0')
(u'c', u'/device:CPU:0')
(u'd', u'/device:GPU:0')

因此，无论创建c时GPU的活动设备上下文如何，它都将tensorc放置在a所在的同一设备上。这对于多GPU训练非常重要。想象一下，如果你不小心，把一个张量相互依赖的图随机放在8个设备上。完全的灾难效率。^{}可以确保这不会发生。

文档中没有对它进行解释，因为它只供内部库使用，所以不能保证它会保留下来。（但很有可能。如果您想了解更多信息，可以从2018年5月起在source code中查找；可能会随着代码的更改而移动。）

你不太可能需要这个，除非你在做一些低级的工作。大多数人只使用一个GPU，即使您使用多个，您通常一次构建一个GPU的图形，即一次在一个^{}上下文管理器中。

使用它的一个例子是^{}类。显然，确保他们正在跟踪是个好主意。