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 перед их сбором.

Вы можете передать один тензор в всеобщий сбор:

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)>)

Вы также можете передать вложенную структуру тензоров в всеобщий сбор, например, список:

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)
([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)>
}],)

Что, если вы собираете тензоры с разными формами на разных репликах? Рассмотрим следующий пример с двумя репликами, где у вас 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))
([PerReplica:{
  0: <tf.Tensor: shape=(), dtype=float32, numpy=2.0>,
  1: <tf.Tensor: shape=(), dtype=float32, numpy=2.0>
}, PerReplica:{
  0: <tf.Tensor: shape=(), dtype=float32, numpy=4.0>,
  1: <tf.Tensor: shape=(), dtype=float32, numpy=4.0>
}],)

Обратите внимание, что все реплики должны участвовать во всеобщем уменьшении, в противном случае операция зависает. Обратите внимание, что при наличии нескольких всеобщих уменьшений они должны выполняться в том же порядке на всех репликах. Динамическое распределение всеобщего уменьшения на основе условий обычно чревато ошибками.

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

• tf.distribute.StrategyExtended.reduce_to: API для уменьшения и всеобщего уменьшения в контексте между репликами.
• tf.distribute.StrategyExtended.batch_reduce_to: API для пакетного уменьшения и всеобщего уменьшения в контексте между репликами.
• 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)