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, options=None
)

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

Аналогично reduce_to, но принимает список пар (значение, пункты назначения). Это более эффективно, чем снижение каждого значения по отдельности.

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

• tf.distribute.StrategyExtended.reduce_to: не-пакетная версия этого API.
• tf.distribute.ReplicaContext.all_reduce: аналог этого API в контексте реплики. Он поддерживает как пакетное, так и непакетное все-reduce.
• tf.distribute.Strategy.reduce: более удобный метод для снижения до хоста в контексте межрепликации.

См. reduce_to для получения дополнительной информации.

@tf.function
def step_fn(var):

  def merge_fn(strategy, value, var):
    # All-reduce the value. Note that `value` here is a
    # `tf.distribute.DistributedValues`.
    reduced = strategy.extended.batch_reduce_to(
        tf.distribute.ReduceOp.SUM, [(value, var)])[0]
    strategy.extended.update(var, lambda var, value: var.assign(value),
        args=(reduced,))

  value = tf.identity(1.)
  tf.distribute.get_replica_context().merge_call(merge_fn,
    args=(value, var))

def run(strategy):
  with strategy.scope():
    v = tf.Variable(0.)
    strategy.run(step_fn, args=(v,))
    return v

run(tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]))
MirroredVariable:{
  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>,
  1: <tf.Variable 'Variable/replica_1:0' shape=() dtype=float32, numpy=2.0>
}
run(tf.distribute.experimental.CentralStorageStrategy(
    compute_devices=["GPU:0", "GPU:1"], parameter_device="CPU:0"))
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>
run(tf.distribute.OneDeviceStrategy("GPU:0"))
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>

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, options=None
)

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

reduce_to агрегирует tf.distribute.DistributedValues и распределённые переменные. Он поддерживает как плотные значения, так и tf.IndexedSlices.

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

• tf.distribute.StrategyExtended.batch_reduce_to: пакетная версия этого API.
• tf.distribute.ReplicaContext.all_reduce: аналог этого API в контексте реплики. Он поддерживает как пакетное, так и непакетное все-reduce.
• tf.distribute.Strategy.reduce: более удобный метод для снижения до хоста в межреплицированном контексте.

destinations указывает, куда выполнить снижение значения, например, "GPU:0". Вы также можете передать Tensor, и пункты назначения будут устройством этого тензора. Для всех-reduce передайте то же самое в value и destinations.

Он может быть использован в tf.distribute.ReplicaContext.merge_call для написания кода, который работает для всех tf.distribute.Strategy.

@tf.function
def step_fn(var):

  def merge_fn(strategy, value, var):
    # All-reduce the value. Note that `value` here is a
    # `tf.distribute.DistributedValues`.
    reduced = strategy.extended.reduce_to(tf.distribute.ReduceOp.SUM,
        value, destinations=var)
    strategy.extended.update(var, lambda var, value: var.assign(value),
        args=(reduced,))

  value = tf.identity(1.)
  tf.distribute.get_replica_context().merge_call(merge_fn,
    args=(value, var))

def run(strategy):
  with strategy.scope():
    v = tf.Variable(0.)
    strategy.run(step_fn, args=(v,))
    return v

run(tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]))
MirroredVariable:{
  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>,
  1: <tf.Variable 'Variable/replica_1:0' shape=() dtype=float32, numpy=2.0>
}
run(tf.distribute.experimental.CentralStorageStrategy(
    compute_devices=["GPU:0", "GPU:1"], parameter_device="CPU:0"))
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>
run(tf.distribute.OneDeviceStrategy("GPU:0"))
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>

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