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`	Целое число, `tf.function` или None. Предпочтительнее использовать целое число, чтобы избежать проблем с вложенными `tf.function`. Принимает `tf.function` только для совместимости с `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. Примечание: Это не гарантирует, что это будет тот же ID, что и XLA ID реплики, используемый для операций низкого уровня, таких как 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`	0-мерный тензор int32. Измерение, по которому выполняется сборка.
`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>])

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

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

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

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

Это позволяет обеспечить взаимодействие и координацию при нескольких вызовах функции 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)

Аргументы
`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