tf.distribute.StrategyExtended

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

tf.distribute.StrategyExtended(
    container_strategy
)

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

Некоторые распространённые случаи использования функций на этой странице:

• Локальность

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

• Как обновить распределённую переменную

Распределённая переменная — это переменные, созданные на нескольких устройствах. Как обсуждалось в словаре терминов, зеркальные переменные и переменные SyncOnRead являются двумя примерами. Стандартный шаблон для обновления распределённых переменных заключается в следующем:

1. В своей функции, переданной в tf.distribute.Strategy.run, вычислите список пар (обновление, переменная). Например, обновление может быть градиентом потери относительно переменной.
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 на переменной так же, как вы делаете с обычной tf.Variable. Вы можете вызвать метод как в контексте реплики, так и в контексте кросс-реплики. Для зеркальной переменной, вызов assign в контексте реплики требует от вас указания типа aggregation в конструкторе переменной. В этом случае переключение контекста и синхронизация, описанные в шагах с 2 по 4, обрабатываются за вас. Если вы вызываете assign на зеркальной переменной в контексте кросс-реплики, вы можете назначить только одно значение или значения из другой зеркальной переменной или зеркального tf.distribute.DistributedValues. Для переменной SyncOnRead, в контексте реплики, вы можете просто вызвать assign на ней, и агрегация не происходит в фоновом режиме. В контексте кросс-реплики вы можете назначить только одно значение переменной SyncOnRead. Пример использования — восстановление из контрольной точки: если тип aggregation переменной — tf.VariableAggregation.SUM, предполагается, что значения реплики были сложены до создания контрольной точки, поэтому при восстановлении значение делится на количество реплик, а затем назначается каждой реплике; если тип aggregation — tf.VariableAggregation.MEAN, значение назначается каждой реплике напрямую.

Методы

batch_reduce_to

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

batch_reduce_to(
    reduce_op, value_destination_pairs, experimental_hints=None
)

Объединяет несколько вызовов 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))

reduce_to

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

reduce_to(
    reduce_op, value, destinations, experimental_hints=None
)

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

update

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

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

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

tf.distribute.StrategyExtended.update принимает распределённую переменную var для обновления, функцию обновления fn, и args и kwargs для fn. Она применяет fn к каждой составляющей переменной var и передаёт соответствующие значения из args и kwargs. Ни args, ни kwargs не могут содержать значения, приходящиеся на одну реплику. Если они содержат зеркальные значения, они будут распакованы перед вызовом fn. Например, fn может быть assign_add, а args — зеркальным DistributedValues, где каждый компонент содержит значение, которое нужно добавить к этой зеркальной переменной var. Вызов update вызовет assign_add для каждой составляющей переменной var с соответствующим тензорным значением на этом устройстве.

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

strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1']) # With 2 devices
with strategy.scope():
  v = tf.Variable(5.0, aggregation=tf.VariableAggregation.SUM)
def update_fn(v):
  return v.assign(1.0)
result = strategy.extended.update(v, update_fn)
# result is
# Mirrored:{
#  0: tf.Tensor(1.0, shape=(), dtype=float32),
#  1: tf.Tensor(1.0, shape=(), dtype=float32)
# }

Если 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.

value_container

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

value_container(
    value
)

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

variable_created_in_scope

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

variable_created_in_scope(
    v
)

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

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

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

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

strategy = tf.distribute.MirroredStrategy()
v = tf.Variable(1.)
strategy.extended.variable_created_in_scope(v)
False