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]]}```