Python tf.compat.v1.distribute.ReplicaContext.all_reduce用法及代碼示例

用法

all_reduce(
    reduce_op, value, options=None
)

All-reduces value 跨所有副本。

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
def step_fn():
  ctx = tf.distribute.get_replica_context()
  value = tf.identity(1.)
  return ctx.all_reduce(tf.distribute.ReduceOp.SUM, value)
strategy.experimental_local_results(strategy.run(step_fn))
(<tf.Tensor:shape=(), dtype=float32, numpy=2.0>,
 <tf.Tensor:shape=(), dtype=float32, numpy=2.0>)

它支持批處理操作。您可以傳遞一個值列表，並在可能的情況下嘗試對它們進行批處理。您還可以指定 options 以指示所需的批處理行為，例如將值批處理到多個包中，以便它們可以更好地與計算重疊。

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
def step_fn():
  ctx = tf.distribute.get_replica_context()
  value1 = tf.identity(1.)
  value2 = tf.identity(2.)
  return ctx.all_reduce(tf.distribute.ReduceOp.SUM, [value1, value2])
strategy.experimental_local_results(strategy.run(step_fn))
([<tf.Tensor:shape=(), dtype=float32, numpy=2.0>,
<tf.Tensor:shape=(), dtype=float32, numpy=4.0>],
[<tf.Tensor:shape=(), dtype=float32, numpy=2.0>,
<tf.Tensor:shape=(), dtype=float32, numpy=4.0>])

注意所有副本都需要參與all-reduce，否則這個操作會掛起。請注意，如果有多個all-reduces，它們需要在所有副本上以相同的順序執行。根據條件調度all-reduce通常是error-prone。

已知限製：如果 value 包含 tf.IndexedSlices ，嘗試計算梯度 w.r.t value 將導致錯誤。

此 API 目前隻能在副本上下文中調用。其他減少副本值的變體是：

• tf.distribute.StrategyExtended.reduce_to ：cross-replica 上下文中的 reduce 和 all-reduce API。
• tf.distribute.StrategyExtended.batch_reduce_to ：cross-replica 上下文中的批量歸約和all-reduce API。
• tf.distribute.Strategy.reduce ：在cross-replica 上下文中減少到主機的更方便的方法。