Python tf.distribute.ReplicaContext.all_gather用法及代码示例

用法

all_gather(
    value, axis, options=None
)

All-gathers value 沿 axis 的所有副本。

注意：all_gather 方法只能在副本上下文中调用。对于 cross-replica 上下文对应项，请参阅 tf.distribute.Strategy.gather 。所有副本都需要参与all-gather，否则这个操作会挂起。因此，如果在任何副本中调用all_gather，则必须在所有副本中调用它。

注意：如果有多个 all_gather 调用，则需要在所有副本上以相同的顺序执行它们。根据条件调度all_gather通常是error-prone。

对于除 tf.distribute.TPUStrategy 之外的所有策略，不同副本上的输入 value 必须具有相同的等级，并且它们的形状必须在除 axis -th 维度之外的所有维度上相同。换句话说，它们的形状在d 维度上不能不同，其中d 不等于axis 参数。例如，给定一个 tf.distribute.DistributedValues 在两个副本上具有形状为 (1, 2, 3) 和 (1, 3, 3) 的分量张量，您可以在其上调用 all_gather(..., axis=1, ...) ，但不能调用 all_gather(..., axis=0, ...) 或 all_gather(..., axis=2, ...) 。但是，对于 tf.distribute.TPUStrategy ，所有张量必须具有完全相同的秩和形状。

注意：输入 value 必须具有非零等级。否则，请考虑在收集它们之前使用tf.expand_dims。

您可以将单个张量传递给all-gather：

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
@tf.function
def gather_value():
  ctx = tf.distribute.get_replica_context()
  local_value = tf.constant([1, 2, 3])
  return ctx.all_gather(local_value, axis=0)
result = strategy.run(gather_value)
result
PerReplica:{
  0:<tf.Tensor:shape=(6,), dtype=int32, numpy=array([1, 2, 3, 1, 2, 3], dtype=int32)>,
  1:<tf.Tensor:shape=(6,), dtype=int32, numpy=array([1, 2, 3, 1, 2, 3], dtype=int32)>
}
strategy.experimental_local_results(result)
(<tf.Tensor:shape=(6,), dtype=int32, numpy=array([1, 2, 3, 1, 2, 3],
dtype=int32)>,
<tf.Tensor:shape=(6,), dtype=int32, numpy=array([1, 2, 3, 1, 2, 3],
dtype=int32)>)

您还可以将张量的嵌套结构传递给all-gather，例如一个列表：

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
@tf.function
def gather_nest():
  ctx = tf.distribute.get_replica_context()
  value_1 = tf.constant([1, 2, 3])
  value_2 = tf.constant([[1, 2], [3, 4]])
  # all_gather a nest of `tf.distribute.DistributedValues`
  return ctx.all_gather([value_1, value_2], axis=0)
result = strategy.run(gather_nest)
result
[PerReplica:{
  0:<tf.Tensor:shape=(6,), dtype=int32, numpy=array([1, 2, 3, 1, 2, 3], dtype=int32)>,
  1:<tf.Tensor:shape=(6,), dtype=int32, numpy=array([1, 2, 3, 1, 2, 3], dtype=int32)>
}, PerReplica:{
  0:<tf.Tensor:shape=(4, 2), dtype=int32, numpy=
array([[1, 2],
       [3, 4],
       [1, 2],
       [3, 4]], dtype=int32)>,
  1:<tf.Tensor:shape=(4, 2), dtype=int32, numpy=
array([[1, 2],
       [3, 4],
       [1, 2],
       [3, 4]], dtype=int32)>
}]
strategy.experimental_local_results(result)
([<tf.Tensor:shape=(6,), dtype=int32, numpy=array([1, 2, 3, 1, 2, 3], dtype=int32)>,
<tf.Tensor:shape=(4, 2), dtype=int32, numpy=
array([[1, 2],
       [3, 4],
       [1, 2],
       [3, 4]], dtype=int32)>],
       [<tf.Tensor:shape=(6,), dtype=int32, numpy=array([1, 2, 3, 1, 2, 3], dtype=int32)>,
       <tf.Tensor:shape=(4, 2), dtype=int32, numpy=
array([[1, 2],
       [3, 4],
       [1, 2],
       [3, 4]], dtype=int32)>])

如果你是在不同副本上具有不同形状的all-gathering 张量怎么办？考虑以下具有两个副本的示例，其中 value 作为嵌套结构，由 all-gather、a 和 b 两个项目组成。

在副本 0 上，value 是 {'a':[0], 'b':[[0, 1]]} 。

在副本 1 上，value 是 {'a':[1], 'b':[[2, 3], [4, 5]]} 。

all_gather 与 axis =0 的结果(在每个副本上)是：

```{'a':[1, 2], 'b':[[0, 1], [2, 3], [4, 5]]}```