tf.compat.v2.distribute.StrategyExtended

Дополнительные API для алгоритмов, которым необходимо учитывать распределение.

tf.compat.v2.distribute.StrategyExtended(
    container_strategy
)

Примечание: Для большинства случаев использования tf.distribute.Strategy, вызов этих методов не требуется, так как библиотеки TensorFlow (например, оптимизаторы) уже вызывают эти методы при необходимости от вашего имени.

Концепции более низкого уровня:

• Обёртки значений: Для представления значений параллельно на устройствах (либо реплики, либо устройства, связанные с конкретным значением), мы оборачиваем их в объект "PerReplica" или "Mirrored", который содержит отображение идентификатора реплики к значениям. "PerReplica" используется, когда значение может отличаться между репликами, а "Mirrored" — когда значения одинаковы.
• Разворачивание и слияние: Рассмотрим вызов функции fn на нескольких репликах, например, experimental_run_v2(fn, args=[w]) с аргументом w , который является обёрнутым значением. Это означает, что w будет иметь отображение, сопоставляющее идентификатор реплики 0 со значением w0, идентификатор реплики 11 со значением w1 и т.д. experimental_run_v2() разворачивает w перед вызовом fn, поэтому она вызывает fn(w0) для d0, fn(w1) для d1 и т.д. Затем она объединяет возвращаемые значения из fn(), что может привести к обёрнутым значениям. Например, предположим, что fn() возвращает кортеж с тремя компонентами: (x, a, v0) от реплики 0, (x, b, v1) от реплики 1 и т.д. Если первый компонент — это один и тот же объект x для каждой реплики, то первый компонент объединённого результата также будет x. Если второй компонент различается (a, b, ...) для каждой реплики, то объединённое значение будет иметь отображение, сопоставляющее устройство реплики с различными значениями. Если третий компонент — это члены зеркальной переменной (v отображает d0 на v0, d1 на v1 и т.д.), то объединённый результат будет этой зеркальной переменной (v).
• Устройства рабочих узлов против устройств параметров: Большинство вычислений реплик происходят на устройствах рабочих узлов. Поскольку мы пока не поддерживаем параллелизм моделей, будет по одному устройству рабочего узла на реплику. При использовании параметрических серверов или централизованного хранилища, набор устройств, содержащих переменные, может отличаться; в противном случае устройства параметров могут совпадать с устройствами рабочих узлов.

Контекст реплики против контекста между репликами

Контекст реплики — это когда мы находимся в некоторой функции, которая вызывается один раз для каждой реплики. В противном случае мы находимся в контексте между репликами, который полезен для вызова методов tf.distribute.Strategy, которые работают с репликами (например, reduce_to()). По умолчанию вы начинаете в контексте реплики (в "контексте единственной реплики по умолчанию"), и некоторые методы могут переключать вас туда и обратно. Есть третий режим, в котором вы можете находиться, — контекст обновления, используемый при обновлении переменных.

• tf.distribute.Strategy.scope: переходит в контекст между репликами, когда ни одна другая стратегия не находится в области видимости.
• tf.distribute.Strategy.experimental_run_v2: вызывает функцию в контексте реплики.
• tf.distribute.ReplicaContext.merge_call: переходит из контекста реплики в контекст между репликами.
• tf.distribute.StrategyExtended.update: вызывает функцию в контексте обновления из контекста между репликами.

В контексте реплики вы можете свободно читать значения переменных, но можете обновить их значение только в том случае, если они указывают способ агрегирования обновления с помощью параметра aggregation в конструкторе переменной. В контексте между репликами вы можете читать или писать переменные (записи могут потребоваться транслировать во все копии переменной, если она зеркальная).

Синхронизация при чтении переменных

В некоторых случаях, например, для метрики, мы хотим накапливать множество обновлений на каждой реплике независимо и агрегировать их только при чтении. Это может значительно повысить производительность, когда значение читается только редко (возможно, значение читается только в конце эпохи или при создании контрольной точки). Это переменные, созданные путём передачи synchronization=ON_READ в конструктор переменной (и некоторое значение для aggregation).

Стратегия может выбрать размещение переменной на нескольких устройствах, подобно зеркальным переменным, но в отличие от зеркальных переменных мы не синхронизируем обновления для обеспечения одинакового значения. Вместо этого синхронизация выполняется при чтении в контексте между репликами. В контексте реплики чтение и запись выполняется в локальной копии (мы разрешаем чтение, чтобы вы могли писать код, подобный v = 0.9*v + 0.1*update). Мы не разрешаем операции, подобные v.assign_add в контексте между репликами для переменных с синхронизацией при чтении; в настоящее время для таких обновлений у нас нет случая использования, и в зависимости от режима агрегирования такие обновления могут быть нецелесообразны.

Локальность

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

Значения "Per-replica" существуют на устройствах рабочих узлов с различными значениями для каждой реплики. Они создаются путём итерации по "распределённому Dataset", возвращаемому tf.distribute.Strategy.experimental_distribute_dataset и tf.distribute.Strategy.experimental_distribute_datasets_from_function. Они также являются типичным результатом, возвращаемым tf.distribute.Strategy.experimental_run_v2. Обычно вы не можете напрямую использовать значение per-replica в контексте между репликами, не определив предварительно, как агрегировать значения по репликам, например, используя tf.distribute.Strategy.reduce.

Значения "Mirrored" похожи на значения per-replica, за исключением того, что нам известно, что значение на всех репликах одинаковое. Мы можем безопасно читать зеркальное значение в контексте между репликами, используя значение на любой реплике. Вы можете преобразовать значение per-replica в зеркальное значение, используя tf.distribute.ReplicaContext.all_reduce.

Значения также могут иметь такую же локальность, как и переменная, что представляет собой зеркальное значение, но размещённое на тех же устройствах, что и переменная (в отличие от вычислительных устройств). Такие значения могут быть переданы в вызов tf.distribute.StrategyExtended.update для обновления значения переменной. Вы можете использовать tf.distribute.StrategyExtended.colocate_vars_with для присвоения переменной той же локальности, что и другой переменной. Это полезно, например, для "слотовых" переменных, используемых оптимизатором для отслеживания статистических данных, используемых для обновления основной/модельной переменной. Вы можете преобразовать значение per-replica в локальность переменной, используя tf.distribute.StrategyExtended.reduce_to или tf.distribute.StrategyExtended.batch_reduce_to.

Помимо слотовых переменных, которые должны быть со-размещёны со своими основными переменными, оптимизаторы также определяют не-слотовые переменные. Это могут быть такие вещи, как "количество шагов обновления", выполненных, или "beta1^t" и "beta2^t". У каждой стратегии есть некоторая политика, касающаяся устройств, на которые должны быть скопированы эти переменные, которая называется "устройствами не-слотовых переменных" (некоторая подмножество устройств параметров). Мы требуем, чтобы все не-слотовые переменные были распределены на одном устройстве или были зеркально отображены на том же наборе устройств. Вы можете использовать tf.distribute.StrategyExtended.non_slot_devices для выбора согласованного набора устройств, которые нужно передать как в tf.distribute.StrategyExtended.colocate_vars_with, так и в tf.distribute.StrategyExtended.update_non_slot.

Как обновить переменную

Стандартный шаблон для обновления переменных:

1. В вашей функции, переданной в tf.distribute.Strategy.experimental_run_v2, вычислите список пар (обновление, переменная). Например, обновление может быть градиентом функции потерь по отношению к переменной.
2. Переключитесь в режим между репликами, вызвав tf.distribute.get_replica_context().merge_call() с обновлениями и переменными в качестве аргументов.
3. Вызовите tf.distribute.StrategyExtended.reduce_to(VariableAggregation.SUM, t, v) (для одной переменной) или tf.distribute.StrategyExtended.batch_reduce_to (для списка переменных) для суммирования обновлений и трансляции результата на устройства переменных.
4. Вызовите tf.distribute.StrategyExtended.update(v) для каждой переменной, чтобы обновить её значение.

Шаги 2—4 выполняются автоматически классом tf.keras.optimizers.Optimizer, если вы вызываете его метод tf.keras.optimizers.Optimizer.apply_gradients в контексте реплики. Они также выполняются автоматически, если вы вызываете метод assign* для переменной (не синхронизированной при чтении), которая была создана с методом агрегирования (который используется для определения сокращения, используемого на шаге 3).

Слои, учитывающие распределение

Слои обычно вызываются в контексте реплики, за исключением определения функциональной модели. tf.distribute.in_cross_replica_context позволит определить, в каком случае вы находитесь. Если вы находитесь в контексте реплики, функция tf.distribute.get_replica_context вернёт объект tf.distribute.ReplicaContext. Объект ReplicaContext имеет метод all_reduce для агрегирования по всем репликам. В качестве альтернативы вы можете обновить переменные, выполнив шаги 2-4 выше.

Примечание: Для новых реализаций tf.distribute.Strategy, пожалуйста, поместите весь код в подкласс tf.distribute.StrategyExtended. Единственный код, необходимый для подкласса tf.distribute.Strategy, заключается в создании экземпляра подкласса tf.distribute.StrategyExtended в методе __init__.

Методы

batch_reduce_to

Просмотр исходного кода

batch_reduce_to(
    reduce_op, value_destination_pairs
)

Объедините несколько reduce_to вызовов в один для более быстрого выполнения.

colocate_vars_with

Просмотр исходного кода

colocate_vars_with(
    colocate_with_variable
)

Область, которая управляет тем, на каких устройствах будут созданы переменные.

Внутри этой области не следует добавлять операции в граф, она должна использоваться только при создании переменных (некоторые реализации работают путём изменения создания переменных, другие — путём использования области tf.compat.v1.colocate_with()).

Это может быть использовано только внутри self.scope().

Пример использования:

with strategy.scope():
  var1 = tf.Variable(...)
  with strategy.extended.colocate_vars_with(var1):
    # var2 and var3 will be created on the same device(s) as var1
    var2 = tf.Variable(...)
    var3 = tf.Variable(...)

  def fn(v1, v2, v3):
    # operates on v1 from var1, v2 from var2, and v3 from var3

  # `fn` runs on every device `var1` is on, `var2` and `var3` will be there
  # too.
  strategy.extended.update(var1, fn, args=(var2, var3))

non_slot_devices

Просмотр исходного кода

non_slot_devices(
    var_list
)

Устройство(я) для переменных без слотов.

Создайте переменные на этих устройствах в блоке with colocate_vars_with(non_slot_devices(...)):. Обновляйте их, используя update_non_slot().

reduce_to

Просмотр исходного кода

reduce_to(
    reduce_op, value, destinations
)

Объединить (например, сумму или среднее значение) значения по репликам.

update

Просмотр исходного кода

update(
    var, fn, args=(), kwargs=None, group=True
)

Выполните fn для обновления var с использованием входных данных, отражённых на тех же устройствах.

Если var отражён на нескольких устройствах, тогда это реализует логику типа:

results = {}
for device, v in var:
  with tf.device(device):
    # args and kwargs will be unwrapped if they are mirrored.
    results[device] = fn(v, *args, **kwargs)
return merged(results)

В противном случае возвращает fn(var, *args, **kwargs) соположенный с var.

Ни args, ни kwargs не могут содержать значения для каждой реплики. Если они содержат отражённые значения, они будут распакованы перед вызовом fn.

update_non_slot

Просмотр исходного кода

update_non_slot(
    colocate_with, fn, args=(), kwargs=None, group=True
)

Выполняет fn(*args, **kwargs) на устройствах colocate_with.

value_container

Просмотр исходного кода

value_container(
    value
)

Возвращает контейнер, к которому принадлежит эта реплика value.

variable_created_in_scope

Просмотр исходного кода

variable_created_in_scope(
    v
)

Проверяет, была ли v создана во время активности этой области стратегии.

Переменные, созданные внутри области стратегии, "принадлежат" ей:

with strategy.scope():
  v = tf.Variable(1.)
strategy.variable_created_in_scope(v)
True

Переменные, созданные вне стратегии, ей не принадлежат:

v = tf.Variable(1.)
strategy.variable_created_in_scope(v)
False