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>