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 в контексте реплики. Поддерживает пакетное и непакетное 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"))

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 в контексте реплики. Поддерживает как пакетную, так и обычную операцию 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: , 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