tf.distribute.MultiWorkerMirroredStrategy

Стратегия распределения для синхронного обучения на нескольких рабочих узлах.

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

tf.distribute.MultiWorkerMirroredStrategy(
    cluster_resolver=None, communication_options=None
)

Используется в блокнотах

Эта стратегия реализует синхронное распределенное обучение на нескольких рабочих узлах, каждый из которых может иметь несколько GPU. Подобно tf.distribute.MirroredStrategy, она реплицирует все переменные и вычисления на каждом локальном устройстве. Отличие заключается в использовании распределенной реализации коллективных операций (например, all-reduce), позволяющей нескольким рабочим узлам работать совместно.

Необходимо запустить программу на каждом рабочем узле и правильно настроить cluster_resolver. Например, если вы используете tf.distribute.cluster_resolver.TFConfigClusterResolver, каждому рабочему узлу необходимо задать соответствующие task_type и task_id в переменной окружения TF_CONFIG. Пример TF_CONFIG для worker-0 в кластере из двух рабочих узлов:

TF_CONFIG = '{"cluster": {"worker": ["localhost:12345", "localhost:23456"]}, "task": {"type": "worker", "index": 0} }'

Программа выполняется на каждом рабочем узле без изменений. Обратите внимание, что для коллективных операций требуется участие каждого рабочего узла. Все API tf.distribute и не относящиеся к tf.distribute могут использовать коллективные операции внутри себя, например при создании контрольных точек и сохранении, поскольку чтение tf.Variable с tf.VariableSynchronization.ON_READ приводит к выполнению all-reduce для значения. Поэтому рекомендуется запускать одну и ту же программу на каждом рабочем узле. Выбор логики в зависимости от task_type или task_id рабочего узла может привести к ошибкам.

cluster_resolver.num_accelerators() определяет количество GPU, используемых стратегией. Если оно равно нулю, стратегия использует CPU. На всех рабочих узлах должно использоваться одинаковое количество устройств, иначе поведение не определено.

Эта стратегия не предназначена для TPU. Вместо нее используйте tf.distribute.TPUStrategy.

После настройки TF_CONFIG использование этой стратегии аналогично использованию tf.distribute.MirroredStrategy и tf.distribute.TPUStrategy.

strategy = tf.distribute.MultiWorkerMirroredStrategy()

with strategy.scope():
  model = tf.keras.Sequential([
    tf.keras.layers.Dense(2, input_shape=(5,)),
  ])
  optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)

def dataset_fn(ctx):
  x = np.random.random((2, 5)).astype(np.float32)
  y = np.random.randint(2, size=(2, 1))
  dataset = tf.data.Dataset.from_tensor_slices((x, y))
  return dataset.repeat().batch(1, drop_remainder=True)
dist_dataset = strategy.distribute_datasets_from_function(dataset_fn)

model.compile()
model.fit(dist_dataset)

Также можно написать собственный цикл обучения:

@tf.function
def train_step(iterator):

  def step_fn(inputs):
    features, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(features, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits)

    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

  strategy.run(step_fn, args=(next(iterator),))

for _ in range(NUM_STEP):
  train_step(iterator)

Подробное руководство см. в разделе Обучение на нескольких рабочих узлах с Keras.

Сохранение

Сохранять данные и контрольные точки необходимо на всех рабочих узлах, а не только на одном. Это связано с тем, что для переменных с synchronization=ON_READ при сохранении выполняется агрегирование. Чтобы избежать состояний гонки, рекомендуется сохранять данные по разным путям на каждом рабочем узле. Каждый рабочий узел сохраняет одни и те же данные. Примеры см. в руководстве Обучение на нескольких рабочих узлах с Keras.

Известные проблемы

• tf.distribute.cluster_resolver.TFConfigClusterResolver возвращает неверное количество ускорителей. Стратегия использует все доступные GPU, если cluster_resolver — это tf.distribute.cluster_resolver.TFConfigClusterResolver или None.
• В eager-режиме стратегию необходимо создать до вызова любого другого API TensorFlow.

Методы

distribute_datasets_from_function

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

distribute_datasets_from_function(
    dataset_fn, options=None
)

Распределяет экземпляры tf.data.Dataset, созданные вызовами dataset_fn.

Аргумент dataset_fn, который передают пользователи, — это входная функция, принимающая аргумент tf.distribute.InputContext и возвращающая экземпляр tf.data.Dataset. Ожидается, что набор данных, возвращаемый dataset_fn, уже разбит на пакеты размером на реплику (то есть глобальный размер пакета, деленный на количество реплик в синхронизации) и разделен на шарды. tf.distribute.Strategy.distribute_datasets_from_function не разбивает на пакеты и шарды экземпляр tf.data.Dataset, возвращаемый входной функцией. dataset_fn вызывается на устройстве CPU каждого рабочего узла, и каждый из них создает набор данных, из которого каждая реплика на этом рабочем узле извлекает один пакет входных данных (то есть, если у рабочего узла две реплики, на каждом шаге из Dataset будут извлечены два пакета).

Этот метод можно использовать для решения нескольких задач. Во-первых, он позволяет задать собственную логику разбиения на пакеты и шарды. (В отличие от него, tf.distribute.experimental_distribute_dataset выполняет разбиение на пакеты и шарды автоматически.) Например, если experimental_distribute_dataset не может разделить входные файлы на шарды, этот метод можно использовать для разделения набора данных на шарды вручную (избегая медленного резервного поведения в experimental_distribute_dataset). Если набор данных бесконечен, такое разделение на шарды можно выполнить, создав реплики наборов данных, различающиеся только случайным начальным значением.

Функция dataset_fn должна принимать экземпляр tf.distribute.InputContext, предоставляющий информацию о разбиении на пакеты и репликации входных данных.

С помощью свойства element_spec объекта tf.distribute.DistributedDataset, возвращаемого этим API, можно запросить tf.TypeSpec элементов, возвращаемых итератором. Это можно использовать для задания свойства input_signature объекта tf.function. Пример см. в разделе tf.distribute.DistributedDataset.element_spec.

Примечание: При использовании TPUStrategy порядок обработки данных рабочими узлами при использовании tf.distribute.Strategy.experimental_distribute_dataset или tf.distribute.Strategy.distribute_datasets_from_function не гарантируется. Обычно это требуется при использовании tf.distribute для масштабирования прогнозирования. Однако можно добавить индекс для каждого элемента пакета и упорядочить выходные данные соответствующим образом. Пример упорядочивания выходных данных см. в этом фрагменте кода.

Примечание: Преобразования набора данных с состоянием в настоящее время не поддерживаются для tf.distribute.experimental_distribute_dataset или tf.distribute.distribute_datasets_from_function. Любые операции с состоянием, которые могут присутствовать в наборе данных, в настоящее время игнорируются. Например, если в наборе данных есть map_fn, использующий tf.random.uniform для поворота изображения, граф набора данных зависит от состояния (то есть от случайного начального значения) локального компьютера, на котором выполняется процесс Python.

Руководство по использованию и свойствам этого метода см. в руководстве по распределенному вводу данных). Если вас интересует обработка последнего неполного пакета, прочитайте этот раздел.

experimental_distribute_dataset

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

experimental_distribute_dataset(
    dataset, options=None
)

Создает tf.distribute.DistributedDataset из tf.data.Dataset.

Возвращаемый объект tf.distribute.DistributedDataset можно перебирать так же, как обычные наборы данных. ПРИМЕЧАНИЕ. Пользователь не может добавлять преобразования к объекту tf.distribute.DistributedDataset. Можно только создать итератор или изучить tf.TypeSpec создаваемых им данных. Подробнее см. документацию API для tf.distribute.DistributedDataset.

Пример:

global_batch_size = 2
# Passing the devices is optional.
strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
# Create a dataset
dataset = tf.data.Dataset.range(4).batch(global_batch_size)
# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(dataset)
@tf.function
def replica_fn(input):
  return input*2
result = []
# Iterate over the `tf.distribute.DistributedDataset`
for x in dist_dataset:
  # process dataset elements
  result.append(strategy.run(replica_fn, args=(x,)))
print(result)
[PerReplica:{
  0: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([0])>,
  1: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([2])>
}, PerReplica:{
  0: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([4])>,
  1: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([6])>
}]

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

В приведенном выше фрагменте кода dataset разбивается на пакеты с помощью global_batch_size, а вызов experimental_distribute_dataset для него повторно разбивает dataset на пакеты нового размера, равного глобальному размеру пакета, деленному на количество реплик в синхронизации. Перебор выполняется с помощью цикла for в стиле Python. x — это tf.distribute.DistributedValues, содержащий данные для всех реплик, и каждая реплика получает данные нового размера пакета. tf.distribute.Strategy.run отвечает за передачу данных для каждой реплики из x соответствующему replica_fn, выполняемому на каждой реплике.

Разделение на шарды включает автоматическое распределение по нескольким рабочим узлам и внутри каждого рабочего узла. Во-первых, при распределенном обучении на нескольких рабочих узлах (то есть при использовании tf.distribute.experimental.MultiWorkerMirroredStrategy или tf.distribute.TPUStrategy) автоматическое распределение набора данных по рабочим узлам означает, что каждому рабочему узлу назначается подмножество всего набора данных (если задано правильное значение tf.data.experimental.AutoShardPolicy). Это гарантирует, что на каждом шаге каждый рабочий узел обработает глобальный пакет элементов набора данных, не пересекающихся с элементами других рабочих узлов. Можно выбрать различные параметры автоматического распределения с помощью tf.data.experimental.DistributeOptions. Затем данные распределяются внутри каждого рабочего узла между всеми его устройствами (если их больше одного). Это происходит независимо от автоматического распределения по нескольким рабочим узлам.

Примечание: При автоматическом распределении по нескольким рабочим узлам по умолчанию используется режим tf.data.experimental.AutoShardPolicy.AUTO. Этот режим пытается разделить входной набор данных на шарды по файлам, если набор данных создается из наборов данных для чтения (например, tf.data.TFRecordDataset, tf.data.TextLineDataset и т. д.); в противном случае набор данных разделяется по данным, и каждый рабочий узел считывает весь набор данных, обрабатывая только назначенный ему шард. Однако если входных файлов меньше, чем рабочих узлов, рекомендуется отключить автоматическое распределение набора данных по рабочим узлам, установив для tf.data.experimental.DistributeOptions.auto_shard_policy значение tf.data.experimental.AutoShardPolicy.OFF.

По умолчанию этот метод добавляет преобразование предварительной выборки в конец предоставленного пользователем экземпляра tf.data.Dataset. Аргументом преобразования предварительной выборки, то есть buffer_size, является количество реплик в синхронизации.

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

Примечание: При использовании TPUStrategy порядок обработки данных рабочими узлами при использовании tf.distribute.Strategy.experimental_distribute_dataset или tf.distribute.Strategy.distribute_datasets_from_function не гарантируется. Обычно это требуется при использовании tf.distribute для масштабирования прогнозирования. Однако можно добавить индекс для каждого элемента пакета и упорядочить выходные данные соответствующим образом. Пример упорядочивания выходных данных см. в этом фрагменте кода.

Примечание: Преобразования набора данных с состоянием в настоящее время не поддерживаются для tf.distribute.experimental_distribute_dataset или tf.distribute.distribute_datasets_from_function. Любые операции с состоянием, которые могут присутствовать в наборе данных, в настоящее время игнорируются. Например, если в наборе данных есть map_fn, использующий tf.random.uniform для поворота изображения, граф набора данных зависит от состояния (то есть от случайного начального значения) локального компьютера, на котором выполняется процесс Python.

Руководство по использованию и свойствам этого метода см. в руководстве по распределенному вводу данных. Если вас интересует обработка последнего неполного пакета, прочитайте этот раздел.

experimental_distribute_values_from_function

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

experimental_distribute_values_from_function(
    value_fn
)

Создает tf.distribute.DistributedValues из value_fn.

Эта функция создает tf.distribute.DistributedValues для передачи в run, reduce или другие методы, принимающие распределенные значения, когда наборы данных не используются.

Примеры использования:

1. Возврат постоянного значения для каждой реплики:

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   def value_fn(ctx):
     return tf.constant(1.)
   distributed_values = (
       strategy.experimental_distribute_values_from_function(
          value_fn))
   local_result = strategy.experimental_local_results(
       distributed_values)
   local_result
       (<tf.Tensor: shape=(), dtype=float32, numpy=1.0>,
       <tf.Tensor: shape=(), dtype=float32, numpy=1.0>)
       

2. Распределение значений в массиве на основе replica_id:

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   array_value = np.array([3., 2., 1.])
   def value_fn(ctx):
     return array_value[ctx.replica_id_in_sync_group]
   distributed_values = (
       strategy.experimental_distribute_values_from_function(
           value_fn))
   local_result = strategy.experimental_local_results(
       distributed_values)
   local_result
       (3.0, 2.0)
       

3. Задание значений с помощью num_replicas_in_sync:

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   def value_fn(ctx):
     return ctx.num_replicas_in_sync
   distributed_values = (
       strategy.experimental_distribute_values_from_function(
           value_fn))
   local_result = strategy.experimental_local_results(
       distributed_values)
   local_result
       (2, 2)
       

4. Размещение значений на устройствах и их распределение:

   strategy = tf.distribute.TPUStrategy()
   worker_devices = strategy.extended.worker_devices
   multiple_values = []
   for i in range(strategy.num_replicas_in_sync):
     with tf.device(worker_devices[i]):
       multiple_values.append(tf.constant(1.0))
   
   def value_fn(ctx):
     return multiple_values[ctx.replica_id_in_sync_group]
   
   distributed_values = strategy.
     experimental_distribute_values_from_function(
     value_fn)
   

experimental_local_results

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

experimental_local_results(
    value
)

Возвращает список всех локальных значений для отдельных реплик, содержащихся в value.

Примечание: Возвращаются только значения на рабочем узле, на котором запущен этот клиент. При использовании tf.distribute.Strategy, например tf.distribute.experimental.MultiWorkerMirroredStrategy, каждый рабочий узел является отдельным клиентом, и эта функция возвращает только значения, вычисленные на этом рабочем узле.

gather

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

gather(
    value, axis
)

Собирает value со всех реплик вдоль axis на текущем устройстве.

Для объекта tf.distribute.DistributedValues или объекта типа tf.Tensor value этот API собирает и объединяет конкатенацией value со всех реплик вдоль измерения axis. Результат копируется на «текущее» устройство, которым обычно является CPU рабочего узла, где выполняется программа. Для tf.distribute.TPUStrategy это первый хост TPU. Для многоклиентской tf.distribute.MultiWorkerMirroredStrategy это CPU каждого рабочего узла.

Этот API можно вызывать только в контексте между репликами. Аналог для контекста реплики см. в разделе tf.distribute.ReplicaContext.all_gather.

Примечание: Для всех стратегий, кроме tf.distribute.TPUStrategy, входные данные value на разных репликах должны иметь одинаковый ранг, а их формы должны совпадать во всех измерениях, кроме измерения axis. Иными словами, формы не могут различаться в измерении d, если d не равно аргументу axis. Например, если заданы tf.distribute.DistributedValues с тензорами-компонентами форм (1, 2, 3) и (1, 3, 3) на двух репликах, можно вызвать для него gather(..., axis=1, ...), но нельзя gather(..., axis=0, ...) или gather(..., axis=2, ...). Однако для tf.distribute.TPUStrategy.gather все тензоры должны иметь строго одинаковые ранг и форму.

Примечание: В объекте tf.distribute.DistributedValues value тензоры-компоненты должны иметь ненулевой ранг. В противном случае перед сбором данных рекомендуется использовать tf.expand_dims.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
# A DistributedValues with component tensor of shape (2, 1) on each replica
distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(tf.constant([[1], [2]])))
@tf.function
def run():
  return strategy.gather(distributed_values, axis=0)
run()
<tf.Tensor: shape=(4, 1), dtype=int32, numpy=
array([[1],
       [2],
       [1],
       [2]], dtype=int32)>

Другие сочетания показаны в следующем примере:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1", "GPU:2", "GPU:3"])
single_tensor = tf.reshape(tf.range(6), shape=(1,2,3))
distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(single_tensor))
@tf.function
def run(axis):
  return strategy.gather(distributed_values, axis=axis)
axis=0
run(axis)
<tf.Tensor: shape=(4, 2, 3), dtype=int32, numpy=
array([[[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]]], dtype=int32)>
axis=1
run(axis)
<tf.Tensor: shape=(1, 8, 3), dtype=int32, numpy=
array([[[0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5]]], dtype=int32)>
axis=2
run(axis)
<tf.Tensor: shape=(1, 2, 12), dtype=int32, numpy=
array([[[0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2],
        [3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5]]], dtype=int32)>

reduce

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

reduce(
    reduce_op, value, axis
)

Выполняет редукцию value по репликам и возвращает результат на текущем устройстве.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
def step_fn():
  i = tf.distribute.get_replica_context().replica_id_in_sync_group
  return tf.identity(i)

per_replica_result = strategy.run(step_fn)
total = strategy.reduce("SUM", per_replica_result, axis=None)
total
<tf.Tensor: shape=(), dtype=int32, numpy=1>

Чтобы увидеть, как это выглядит при использовании нескольких реплик, рассмотрим тот же пример с MirroredStrategy и двумя GPU:

strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
def step_fn():
  i = tf.distribute.get_replica_context().replica_id_in_sync_group
  return tf.identity(i)

per_replica_result = strategy.run(step_fn)
# Check devices on which per replica result is:
strategy.experimental_local_results(per_replica_result)[0].device
# /job:localhost/replica:0/task:0/device:GPU:0
strategy.experimental_local_results(per_replica_result)[1].device
# /job:localhost/replica:0/task:0/device:GPU:1

total = strategy.reduce("SUM", per_replica_result, axis=None)
# Check device on which reduced result is:
total.device
# /job:localhost/replica:0/task:0/device:CPU:0

Обычно этот API используется для агрегирования результатов, возвращенных разными репликами, например для создания отчетов. Например, значения потерь, вычисленные на разных репликах, можно усреднить с помощью этого API перед выводом.

Примечание: Результат копируется на «текущее» устройство — обычно CPU рабочего узла, где выполняется программа. Для TPUStrategy это первый хост TPU. Для многоклиентской MultiWorkerMirroredStrategy это CPU каждого рабочего узла.

В tf.distribute есть несколько API для редукции значений по репликам:

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

Каким должно быть значение axis?

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

Например, если глобальный размер пакета равен 8, а число реплик равно 2, значения для примеров [0, 1, 2, 3] будут на реплике 0, а [4, 5, 6, 7] — на реплике 1. При значении axis=None, reduce будет выполнять агрегирование только по репликам и возвращать [0+4, 1+5, 2+6, 3+7]. Это полезно, когда каждая реплика вычисляет скалярное или какое-либо другое значение без измерения «пакета», например градиент или потери.

strategy.reduce("sum", per_replica_result, axis=None)

Иногда требуется выполнить агрегирование и по глобальному пакету, и по всем репликам. Для этого укажите измерение пакета в качестве axis, обычно axis=0. В этом случае будет возвращен скаляр 0+1+2+3+4+5+6+7.

strategy.reduce("sum", per_replica_result, axis=0)

Если последний пакет неполный, необходимо указать axis, чтобы результирующая форма была одинаковой для всех реплик. Поэтому, если размер последнего пакета равен 6 и он разделен на [0, 1, 2, 3] и [4, 5], без указания axis=0 возникнет несоответствие форм. Если указать tf.distribute.ReduceOp.MEAN, то при использовании axis=0 будет применен правильный знаменатель, равный 6. Сравните это с вычислением reduce_mean для получения скалярного значения на каждой реплике и последующим усреднением этих средних с помощью этой функции: в таком случае некоторые значения будут иметь вес 1/8, а другие — 1/4.

run

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

run(
    fn, args=(), kwargs=None, options=None
)

Вызывает fn на каждой реплике с заданными аргументами.

Этот метод — основной способ распределения вычислений с помощью объекта tf.distribute. Он вызывает fn на каждой реплике. Если args или kwargs содержат tf.distribute.DistributedValues, например созданные с помощью tf.distribute.DistributedDataset из tf.distribute.Strategy.experimental_distribute_dataset или tf.distribute.Strategy.distribute_datasets_from_function, то при выполнении fn на определённой реплике будет использован компонент tf.distribute.DistributedValues, соответствующий этой реплике.

fn вызывается в контексте реплики. fn может вызвать tf.distribute.get_replica_context(), чтобы получить доступ к таким элементам, как all_reduce. Описание понятия контекста реплики см. в строке документации модуля tf.distribute.

Все аргументы в args или kwargs могут представлять собой вложенную структуру тензоров, например список тензоров. В этом случае args и kwargs будут переданы в fn, вызываемый на каждой реплике. Кроме того, args или kwargs могут быть объектами tf.distribute.DistributedValues, содержащими тензоры или составные тензоры, то есть tf.compat.v1.TensorInfo.CompositeTensor. В этом случае при каждом вызове fn будет получен компонент объекта tf.distribute.DistributedValues, соответствующий данной реплике. Обратите внимание, что произвольные значения Python, не относящиеся к перечисленным выше типам, не поддерживаются.

Пример использования:

1. Входной тензор-константа.

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   tensor_input = tf.constant(3.0)
   @tf.function
   def replica_fn(input):
     return input*2.0
   result = strategy.run(replica_fn, args=(tensor_input,))
   result
       PerReplica:{
         0: <tf.Tensor: shape=(), dtype=float32, numpy=6.0>,
         1: <tf.Tensor: shape=(), dtype=float32, numpy=6.0>
       }
       

2. Входные данные DistributedValues.

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   @tf.function
   def run():
     def value_fn(value_context):
       return value_context.num_replicas_in_sync
     distributed_values = (
       strategy.experimental_distribute_values_from_function(
         value_fn))
     def replica_fn2(input):
       return input*2
     return strategy.run(replica_fn2, args=(distributed_values,))
   result = run()
   result
       <tf.Tensor: shape=(), dtype=int32, numpy=4>
       

3. Используйте tf.distribute.ReplicaContext для выполнения all-reduce значений.

   strategy = tf.distribute.MirroredStrategy(["gpu:0", "gpu:1"])
   @tf.function
   def run():
      def value_fn(value_context):
        return tf.constant(value_context.replica_id_in_sync_group)
      distributed_values = (
          strategy.experimental_distribute_values_from_function(
              value_fn))
      def replica_fn(input):
        return tf.distribute.get_replica_context().all_reduce(
            "sum", input)
      return strategy.run(replica_fn, args=(distributed_values,))
   result = run()
   result
       PerReplica:{
         0: <tf.Tensor: shape=(), dtype=int32, numpy=1>,
         1: <tf.Tensor: shape=(), dtype=int32, numpy=1>
       }
       

scope

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

scope()

Менеджер контекста, который делает стратегию текущей и распределяет переменные.

Этот метод возвращает менеджер контекста и используется следующим образом:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
# Variable created inside scope:
with strategy.scope():
  mirrored_variable = tf.Variable(1.)
mirrored_variable
MirroredVariable:{
  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>,
  1: <tf.Variable 'Variable/replica_1:0' shape=() dtype=float32, numpy=1.0>
}
# Variable created outside scope:
regular_variable = tf.Variable(1.)
regular_variable
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>

Что происходит при входе в Strategy.scope?

• strategy устанавливается в глобальном контексте как «текущая» стратегия. Внутри этой области tf.distribute.get_strategy() будет возвращать эту стратегию. За пределами этой области он возвращает стратегию по умолчанию, не выполняющую никаких действий.
• Вход в область также переводит выполнение в «контекст между репликами». Объяснение контекстов между репликами и реплик см. в разделе tf.distribute.StrategyExtended.
• Создание переменных внутри scope перехватывается стратегией. Каждая стратегия определяет, как именно она влияет на создание переменных. Синхронные стратегии, такие как MirroredStrategy, TPUStrategy и MultiWorkerMiroredStrategy, создают реплицированные переменные на каждой реплике, тогда как ParameterServerStrategy создаёт переменные на серверах параметров. Для этого используется пользовательский tf.variable_creator_scope.
• В некоторых стратегиях также может задаваться область действия устройства по умолчанию: в MultiWorkerMiroredStrategy для каждого рабочего узла задаётся область действия устройства по умолчанию "/CPU:0".

Примечание: Вход в область не приводит к автоматическому распределению вычислений, кроме случаев использования высокоуровневой обучающей платформы, например keras model.fit. Если вы не используете model.fit, необходимо явно распределить вычисления с помощью API strategy.run. Пример см. в руководстве по пользовательскому циклу обучения.

Что должно находиться внутри области, а что — за её пределами?

Существует ряд требований к операциям, которые должны выполняться внутри области. Однако в тех случаях, когда известно, какая стратегия используется, мы часто сами входим в область, поэтому пользователю не нужно делать это явно (то есть вызывать эти операции можно как внутри, так и снаружи области).

• Любая функция, создающая переменные, которые должны быть распределёнными, должна вызываться внутри strategy.scope. Этого можно добиться, вызвав функцию создания переменных непосредственно в контексте области или используя другой API, например strategy.run или keras.Model.fit, который выполнит вход в область автоматически. Переменные, созданные за пределами области, не будут распределены, что может повлиять на производительность. В TF переменные создают, например, модели, оптимизаторы и метрики. Такие объекты всегда следует инициализировать в области, а функции, которые могут создавать переменные отложенно (например, Model.call(), трассировка tf.function и т. д.), также следует вызывать внутри области. Ещё одним источником создания переменных может быть восстановление контрольной точки, если переменные создаются отложенно. Обратите внимание, что любая переменная, созданная внутри стратегии, сохраняет информацию о стратегии. Поэтому чтение и запись таких переменных за пределами strategy.scope также могут выполняться без проблем, и пользователю не нужно входить в область.
• Некоторые API стратегий (например, strategy.run и strategy.reduce), для использования которых требуется находиться в области стратегии, входят в неё автоматически. Поэтому при использовании этих API не нужно явно входить в область самостоятельно.
• Если tf.keras.Model создаётся внутри strategy.scope, объект Model сохраняет информацию об области. При вызове высокоуровневых методов обучающей платформы, таких как model.compile, model.fit и т. д., сохранённая область будет активирована автоматически, а связанная с ней стратегия будет использоваться для распределённого обучения и других операций. Подробный пример см. в руководстве по распределённому обучению Keras. ПРЕДУПРЕЖДЕНИЕ: простой вызов model(..) не активирует автоматически сохранённую область — такое поведение поддерживается только API высокоуровневой обучающей платформы: model.compile, model.fit, model.evaluate, model.predict и model.save можно вызывать как внутри, так и за пределами области.
• Следующие действия можно выполнять как внутри, так и за пределами области:
  • Создание входных наборов данных
  • Определение tf.function, представляющих шаг обучения
  • API сохранения, например tf.saved_model.save. При загрузке создаются переменные, поэтому её следует выполнять внутри области, если вы хотите обучать модель распределённым способом.
  • Сохранение контрольных точек. Как упоминалось выше, checkpoint.restore иногда необходимо выполнять внутри области, если при этом создаются переменные.