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

Аргументы
`strategy`	`tf.distribute.Strategy`.
`replica_id_in_sync_group`	Целое число, `Tensor` или None. По возможности используйте целое число, чтобы избежать проблем с вложенными `tf.function`. Значение `Tensor` поддерживается только для совместимости с `tpu.replicate`.

Атрибуты
`devices`	Возвращает устройства, на которых должна выполняться эта реплика, в виде кортежа строк. (Устарело) Устарело: ЭТА ФУНКЦИЯ УСТАРЕЛА. Она будет удалена в будущей версии. Инструкции по обновлению: не используйте свойство devices. Примечание: Для `tf.distribute.MirroredStrategy` и `tf.distribute.experimental.MultiWorkerMirroredStrategy` возвращается вложенный список строк с названиями устройств, например [["GPU:0"]].
`num_replicas_in_sync`	Возвращает количество реплик, синхронизированных между собой.
`replica_id_in_sync_group`	Возвращает идентификатор реплики. Он идентифицирует реплику среди всех реплик, синхронизированных между собой. Значение идентификатора реплики находится в диапазоне от 0 до `tf.distribute.ReplicaContext.num_replicas_in_sync` - 1. Примечание: Не гарантируется, что это будет тот же идентификатор, что и идентификатор реплики XLA, используемый для низкоуровневых операций, например collective_permute.
`strategy`	Текущий объект `tf.distribute.Strategy`.

Методы

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

Аргументы
`value`	Вложенная структура из объектов `tf.Tensor`, поддерживаемая функцией `tf.nest.flatten`, или экземпляр `tf.distribute.DistributedValues`. Структура объектов `tf.Tensor` должна быть одинаковой во всех репликах. Базовые тензорные объекты могут быть только плотными тензорами ненулевого ранга, но не объектами `tf.IndexedSlices`.
`axis`	Тензор int32 ранга 0. Измерение, вдоль которого выполняется сбор.
`options`	Объект `tf.distribute.experimental.CommunicationOptions`. Параметры выполнения коллективных операций. Они переопределяют параметры по умолчанию, если конструктор объекта `tf.distribute.Strategy` принимает такие параметры. Подробнее о параметрах см. в разделе `tf.distribute.experimental.CommunicationOptions`.

Возвращает
Вложенная структура из объектов `tf.Tensor` с собранными значениями. Структура совпадает со структурой `value`.

`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 в зависимости от условий обычно приводит к ошибкам.

Известное ограничение: если 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: более удобный метод для свертки значений на хосте в контексте между репликами.

Аргументы
`reduce_op`	Значение `tf.distribute.ReduceOp`, задающее способ объединения значений. Можно использовать строковое представление перечисления, например "SUM" или "MEAN".
`value`	Потенциально вложенная структура из объектов `tf.Tensor` или `tf.IndexedSlices`, поддерживаемая функцией `tf.nest.flatten`. Структура и формы `value` должны быть одинаковыми во всех репликах.
`options`	Объект `tf.distribute.experimental.CommunicationOptions`. Параметры выполнения коллективных операций. Они переопределяют параметры по умолчанию, если конструктор объекта `tf.distribute.Strategy` принимает такие параметры. Подробнее о параметрах см. в разделе `tf.distribute.experimental.CommunicationOptions`.

Возвращает
Вложенная структура из объектов `tf.Tensor` со свёрнутыми значениями. Структура совпадает со структурой `value`.

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

Аргументы
`merge_fn`	Функция, объединяющая аргументы из потоков, переданные в виде PerReplica. В качестве первого аргумента принимает объект `tf.distribute.Strategy`.
`args`	Список или кортеж позиционных аргументов для каждого потока, передаваемых в `merge_fn`.
`kwargs`	Словарь именованных аргументов для каждого потока, передаваемых в `merge_fn`.

Возвращает
Возвращаемое значение `merge_fn`, кроме значений `PerReplica`, которые распаковываются.

© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/distribute/ReplicaContext