tf.compat.v1.distribute.ReplicaContext

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

tf.compat.v1.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_reduce

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

all_reduce(
    reduce_op, value, options=None
)

Выполняет all-reduce операцию 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>
}],)

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

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

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