tf.distribute.StrategyExtended

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

Наследуется от: StrategyExtended

tf.distribute.StrategyExtended(
    container_strategy
)

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

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

• Оборачиваемые значения: Чтобы представить значения параллельно на разных устройствах (либо репликах, либо устройствах, связанных с определённым значением), мы оборачиваем их в объект "PerReplica" или "Mirrored", который содержит отображение от id реплики к значениям. "PerReplica" используется, когда значение может отличаться между репликами, а "Mirrored", когда значения одинаковы.
• Разворачивание и слияние: Рассмотрим вызов функции fn на нескольких репликах, например, experimental_run_v2(fn, args=[w]) с аргументом w , являющимся обернутым значением. Это означает, что w будет иметь отображение, принимающее id реплики 0 к w0, id реплики 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 в контексте между репликами для переменных, синхронизируемых при чтении; в настоящее время у нас нет такой необходимости в обновлении, и в зависимости от режима агрегирования такие обновления могут быть нецелесообразны.

Локальность

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

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

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

Значения также могут иметь такую же локальность, как переменная, которая является зеркальным значением, но располагается на тех же устройствах, что и переменная (в отличие от вычислительных устройств). Такие значения могут быть переданы в вызов tf.distribute.StrategyExtended.update для обновления значения переменной. Вы можете использовать tf.distribute.StrategyExtended.colocate_vars_with для присвоения переменной той же локальности, что и другой переменной. Это полезно, например, для переменных "слотов", используемых оптимизатором для отслеживания статистических данных, используемых для обновления основной/модельной переменной. Вы можете преобразовать значение по реплике в локальность переменной, используя 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 в один для более быстрого выполнения.

broadcast_to

Просмотреть исходный код

broadcast_to(
    tensor, destinations
)

Зеркалирование тензора на одном устройстве на всех устройствах рабочих узлов.

call_for_each_replica

Просмотреть исходный код

call_for_each_replica(
    fn, args=(), kwargs=None
)

Выполнение fn один раз на реплике.

fn может вызывать tf.get_replica_context() для доступа к методам, таким как replica_id_in_sync_group и merge_call().

merge_call() используется для взаимодействия между репликами и повторного входа в контекст межреплики. Все реплики приостанавливают выполнение, встретив вызов merge_call(). После этого выполняется функция merge_fn. Результаты затем распаковываются и возвращаются каждому вызову реплики. После этого выполнение возобновляется до завершения fn или встречи с другим вызовом merge_call(). Пример:

# Called once in "cross-replica" context.
def merge_fn(distribution, three_plus_replica_id):
  # sum the values across replicas
  return sum(distribution.experimental_local_results(three_plus_replica_id))

# Called once per replica in `distribution`, in a "replica" context.
def fn(three):
  replica_ctx = tf.get_replica_context()
  v = three + replica_ctx.replica_id_in_sync_group
  # Computes the sum of the `v` values across all replicas.
  s = replica_ctx.merge_call(merge_fn, args=(v,))
  return s + v

with distribution.scope():
  # in "cross-replica" context
  ...
  merged_results = distribution.experimental_run_v2(fn, args=[3])
  # merged_results has the values from every replica execution of `fn`.
  # This statement prints a list:
  print(distribution.experimental_local_results(merged_results))

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

experimental_make_numpy_dataset

Просмотреть исходный код

experimental_make_numpy_dataset(
    numpy_input, session=None
)

Создает набор данных для входных данных, предоставленных с помощью массива NumPy.

Это позволяет избежать добавления numpy_input в качестве большого константы в графе и копирует данные на машину или машины, которые будут обрабатывать входные данные.

experimental_run_steps_on_iterator

Просмотреть исходный код

experimental_run_steps_on_iterator(
    fn, iterator, iterations=1, initial_loop_values=None
)

УСТАРЕЛО: используйте experimental_run_v2 вместо этого.

Выполнение fn с данными из iterator в течение iterations шагов.

Этот метод может использоваться для выполнения функции шага обучения несколько раз с использованием данных из набора данных.

non_slot_devices

Просмотреть исходный код

non_slot_devices(
    var_list
)

Устройства для переменных, не являющихся слотами.

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

read_var

Просмотреть исходный код

read_var(
    v
)

Чтение значения переменной.

Возвращает агрегированное значение реплики-локальной переменной или (только для чтения) значение любой другой переменной.

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