tf.compat.v1.distribute.StrategyExtended

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

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

tf.compat.v1.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 в контексте реплики. Поддерживает как пакетное, так и непакетное выполнение операции all-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: , 1: } run(tf.distribute.experimental.CentralStorageStrategy( compute_devices=["GPU:0", "GPU:1"], parameter_device="CPU:0")) run(tf.distribute.OneDeviceStrategy("GPU:0"))

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.run(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
)

Выполняет fn с входными данными из iterator в течение iterations итераций.

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

non_slot_devices

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

non_slot_devices(
    var_list
)

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

Этот метод возвращает устройства без слотов, на которых размещаются переменные без слотов. Пользователи могут создавать переменные без слотов на этих устройствах, используя блок:

with tf.distribute.StrategyExtended.colocate_vars_with(tf.distribute.StrategyExtended.non_slot_devices(...)):
  ...

read_var

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

read_var(
    v
)

Считывает значение переменной.

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

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; в этом случае целевым устройством будет устройство этого тензора. Для агрегирования всех значений передайте одинаковые значения в 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: , 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.

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 создана при активной области видимости этой стратегии.

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

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