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


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


用法

reduce_to(
    reduce_op, value, destinations, options=None
)

参数

返回

  • 减少到 destinations 的张量或值。

跨副本组合(通过例如总和或平均值)值。

reduce_to 聚合 tf.distribute.DistributedValues 和分布式变量。它支持密集值和 tf.IndexedSlices

此 API 目前只能在cross-replica 上下文中调用。其他减少副本值的变体是:

destinations 指定将值减少到的位置,例如"GPU:0"。您也可以传入 Tensor ,目标将是该张量的设备。对于 all-reduce,将其传递给 valuedestinations

它可用于 tf.distribute.ReplicaContext.merge_call 编写适用于所有 tf.distribute.Strategy 的代码。

@tf.function
def step_fn(var):

  def merge_fn(strategy, value, var):
    # All-reduce the value. Note that `value` here is a
    # `tf.distribute.DistributedValues`.
    reduced = strategy.extended.reduce_to(tf.distribute.ReduceOp.SUM,
        value, destinations=var)
    strategy.extended.update(var, lambda var, value:var.assign(value),
        args=(reduced,))

  value = tf.identity(1.)
  tf.distribute.get_replica_context().merge_call(merge_fn,
    args=(value, var))

def run(strategy):
  with strategy.scope():
    v = tf.Variable(0.)
    strategy.run(step_fn, args=(v,))
    return v

run(tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]))
MirroredVariable:{
  0:<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>,
  1:<tf.Variable 'Variable/replica_1:0' shape=() dtype=float32, numpy=2.0>
}
run(tf.distribute.experimental.CentralStorageStrategy(
    compute_devices=["GPU:0", "GPU:1"], parameter_device="CPU:0"))
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>
run(tf.distribute.OneDeviceStrategy("GPU:0"))
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>

相关用法


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