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

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
)

УСТАНОВЛЕНО: пожалуйста, используйте run вместо этого.

Выполнить fn с входными данными из iterator iterations раз.

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

non_slot_devices

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

non_slot_devices(
    var_list
)

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

УСТАНОВЛЕНО: ТОЛЬКО для TF 1.x.

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

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 в контексте реплики. Поддерживает как пакетный, так и непочтовый all-reduce.
• 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):
    # 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.

update_non_slot

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

update_non_slot(
    colocate_with, fn, args=(), kwargs=None, group=True
)

Выполняет fn(*args, **kwargs) на устройствах colocate_with.

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

УСТАРЕВШАЯ функция: ТОЛЬКО для TF 1.x.

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