当前位置: 首页>>代码示例 >>用法及示例精选 >>正文


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


用法

colocate_vars_with(
    colocate_with_variable
)

参数

  • colocate_with_variable 在此策略的 scope() 中创建的变量。在返回的上下文管理器中创建的变量将与 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))

相关用法


注:本文由纯净天空筛选整理自tensorflow.org大神的英文原创作品 tf.distribute.StrategyExtended.colocate_vars_with。非经特殊声明,原始代码版权归原作者所有,本译文未经允许或授权,请勿转载或复制。