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

用法

batch_reduce_to(
    reduce_op, value_destination_pairs, options=None
)

将多个 reduce_to 调用合并为一个以加快执行速度。

类似于 reduce_to ，但接受 (value, destinations) 对的列表。它比单独减少每个值更有效。

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

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

有关详细信息，请参阅reduce_to。

@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.batch_reduce_to(
        tf.distribute.ReduceOp.SUM, [(value, var)])[0]
    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>