Python tf.distribute.StrategyExtended.colocate_vars_with用法及代码示例

用法

colocate_vars_with(
    colocate_with_variable
)

控制将在其上创建哪些设备变量的范围。

不应在此范围内向图中添加任何操作，仅应在创建变量时使用它(某些实现通过更改变量创建工作，其他实现通过使用 tf.compat.v1.colocate_with() 范围工作)。

这只能在 self.scope() 内部使用。

示例用法：

with strategy.scope():
  var1 = tf.Variable(...)
  with strategy.extended.colocate_vars_with(var1):
    # var2 and var3 will be created on the same device(s) as var1
    var2 = tf.Variable(...)
    var3 = tf.Variable(...)

  def fn(v1, v2, v3):
    # operates on v1 from var1, v2 from var2, and v3 from var3

  # `fn` runs on every device `var1` is on, `var2` and `var3` will be there
  # too.
  strategy.extended.update(var1, fn, args=(var2, var3))