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 上下文中减少到主机的更方便的方法。