tensorflow 中 colocate_with 的用途是啥？

Posted 2023-02-16

【中文标题】tensorflow 中 colocate_with 的用途是啥？

【英文标题】：What is colocate_with used for in tensorflow?tensorflow 中 colocate_with 的用途是什么？

【发布时间】：2017-07-27 03:50:07

【问题描述】：

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

【参考方案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'/设备:CPU:0') (u'b', u'/device:GPU:0') (u'c', u'/设备:CPU:0') (u'd', u'/device:GPU:0')

因此它将张量 c 放置在 a 所在的同一设备上，而不管活动设备上下文如何创建 c 时的 GPU。这对于多 GPU 训练非常重要。想象一下，如果您不小心，并且将张量相互依赖的图随机放置在 8 个设备上。一个完整的灾难效率明智的。 tf.colocate_with() 可以确保不会发生这种情况。

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

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

使用它的一个例子是tf.train.ExponentialMovingAverage 类。显然，确保他们正在跟踪colocate the decay and moving average variables with the value tensor 看起来是个好主意。