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))