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