tf.distribute.ReplicaContext

Класс с набором API-функций, которые можно вызвать в контексте реплики.

tf.distribute.ReplicaContext(
    strategy, replica_id_in_sync_group
)

Вы можете использовать tf.distribute.get_replica_context, чтобы получить экземпляр ReplicaContext, который может быть вызван только внутри функции, переданной в tf.distribute.Strategy.run.

strategy = tf.distribute.MirroredStrategy(['GPU:0', 'GPU:1'])
def func():
  replica_context = tf.distribute.get_replica_context()
  return replica_context.replica_id_in_sync_group
strategy.run(func)
PerReplica:{
  0: <tf.Tensor: shape=(), dtype=int32, numpy=0>,
  1: <tf.Tensor: shape=(), dtype=int32, numpy=1>
}

Методы

all_gather

Просмотреть исходный код

all_gather(
    value, axis, options=None
)

Всеобщая сборка value по всем репликам вдоль axis.

Примечание: Метод all_gather может быть вызван только в контексте реплики. Для аналога в контексте кросс-реплик см. tf.distribute.Strategy.gather. Все реплики должны участвовать в всеобщей сборке, иначе эта операция зависнет. Поэтому, если all_gather вызывается в любой реплике, она должна быть вызвана во всех репликах.

Примечание: Если вызывается несколько all_gather операций, они должны выполняться в том же порядке на всех репликах. Вызов all_gather на основе условий обычно приводит к ошибкам.

Для всех стратегий, кроме tf.distribute.TPUStrategy, вход value на разных репликах должен иметь одинаковый ранг, и их формы должны быть одинаковыми во всех измерениях, кроме axis-го измерения. Другими словами, их формы не могут отличаться в измерении 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)>])

Что делать, если вы собираете тензоры с разными формами на разных репликах? Рассмотрим следующий пример с двумя репликами, где у вас есть value как вложенная структура, состоящая из двух элементов для сборки, 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]]}
  

all_reduce

Просмотреть исходный код

all_reduce(
    reduce_op, value, options=None
)

Всеобщий расчет 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-reduce операций, они должны выполняться в том же порядке на всех репликах. Вызов all-reduce на основе условий обычно приводит к ошибкам.

Известное ограничение: если value содержит tf.IndexedSlices, попытка вычислить градиент по value приведет к ошибке.

Этот API в данный момент может быть вызван только в контексте реплики. Другие варианты уменьшения значений по репликам:

• tf.distribute.StrategyExtended.reduce_to: API для уменьшения и all-reduce в контексте кросс-реплик.
• tf.distribute.StrategyExtended.batch_reduce_to: API для пакетного уменьшения и all-reduce в контексте кросс-реплик.
• tf.distribute.Strategy.reduce: более удобный метод для уменьшения до хоста в контексте кросс-реплик.

merge_call

Просмотреть исходный код

merge_call(
    merge_fn, args=(), kwargs=None
)

Объединяет аргументы по репликам и выполняет merge_fn в контексте кросс-реплик.

Это позволяет выполнять коммуникацию и координацию, когда имеется несколько вызовов step_fn, вызванных вызовом strategy.run(step_fn, ...).

См. tf.distribute.Strategy.run для объяснения.

Если не внутри распределённого контекста, это эквивалентно:

strategy = tf.distribute.get_strategy()
with cross-replica-context(strategy):
  return merge_fn(strategy, *args, **kwargs)