什么是colocate_with用于张量流？

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',你'/设备:GPU:0')

因此,无论在创建c时GPU的活动设备上下文如何,它都将张量c置于a所在的同一设备上.这对于多GPU培训非常重要.想象一下,如果你不小心,你有一个图表,其中张量依赖于彼此随机放置在8个设备上.一个完整的灾难效率方面.可以确保不会发生这种情况.tf.colocate_with()

它没有在文档中解释,因为它只是由内部库使用,所以不保证它会保留.(但很可能会这样.如果你想了解更多,你可以在2018年5月的源代码中查找它;可能会随着代码的更改而移动.)

除非你正在处理一些低级别的东西,否则你不太可能需要它.大多数人只使用一个GPU,即使你使用多个GPU,你通常一次只能在一个GPU中构建一个GPU,一次只能在一个tf.device()上下文管理器中.

使用它的一个例子是tf.train.ExponentialMovingAverage类.显然,确保将衰变和移动平均变量与它们正在跟踪的值张量相关联是一个好主意.