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))
([<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 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 в контексте между репликами.

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

Пояснение см. в разделе tf.distribute.Strategy.run.

Если вызов выполняется вне распределённой области, он эквивалентен следующему:

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