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

用法

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

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

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

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

• tf.distribute.StrategyExtended.batch_reduce_to ：此 API 的批处理版本。
• tf.distribute.ReplicaContext.all_reduce ：此 API 在副本上下文中的对应项。它支持批处理和非批处理all-reduce。
• tf.distribute.Strategy.reduce ：在cross-replica 上下文中减少到主机的更方便的方法。

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

它可用于 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>