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 в репликационном контексте. Он поддерживает как пакетное, так и непакетное всеобщее уменьшение.
• tf.distribute.Strategy.reduce: более удобный метод для уменьшения до хоста в межрепликационном контексте.

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

@tf.function def step_fn(var):

def merge_fn(strategy, value, var): # Всеобщее уменьшение значения. Обратите внимание, что value здесь — # 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: , 1: } run(tf.distribute.experimental.CentralStorageStrategy( compute_devices=["GPU:0", "GPU:1"], parameter_device="CPU:0")) run(tf.distribute.OneDeviceStrategy("GPU: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 в репликационном контексте. Он поддерживает как пакетное, так и непакетное всеобщее уменьшение.
• tf.distribute.Strategy.reduce: более удобный метод для уменьшения до хоста в межрепликационном контексте.

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

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

@tf.function def step_fn(var):

def merge_fn(strategy, value, var): # Всеreduce значение. Обратите внимание, что value здесь — # 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: , 1: } run(tf.distribute.experimental.CentralStorageStrategy( compute_devices=["GPU:0", "GPU:1"], parameter_device="CPU:0")) run(tf.distribute.OneDeviceStrategy("GPU: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