Python tf.compat.v1.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>