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
)

Выполняет всеобщее сокращение 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>])

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

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

В настоящее время этот 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)