tf.distribute.experimental.MultiWorkerMirroredStrategy

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

Наследует: MultiWorkerMirroredStrategy, Strategy

tf.distribute.experimental.MultiWorkerMirroredStrategy(
    communication=tf.distribute.experimental.CollectiveCommunication.AUTO,
    cluster_resolver=None
)

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

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

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

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

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

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

Эта стратегия не предназначена для 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 возвращает неверное количество ускорителей. Стратегия использует все доступные графические процессоры, если 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 вызывается на устройстве ЦП каждого рабочего узла и создаёт набор данных, из которого каждая реплика на этом рабочем узле извлекает один пакет входных данных (то есть если у рабочего узла две реплики, на каждом шаге из 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. Результат копируется на «текущее» устройство, которым обычно является ЦП рабочего узла, где выполняется программа. Для tf.distribute.TPUStrategy это первый хост TPU. Для многоклиентской tf.distribute.MultiWorkerMirroredStrategy это ЦП каждого рабочего узла.

Этот 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 и двумя графическими процессорами:

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 перед выводом.

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

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

• 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. Пример приведён в руководстве по пользовательскому циклу обучения.

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

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

• Любой код, создающий переменные, которые должны быть распределёнными переменными, необходимо вызывать внутри strategy.scope. Это можно сделать, вызвав функцию создания переменной непосредственно в контексте области или воспользовавшись другим API, например strategy.run или keras.Model.fit, который автоматически выполнит вход в область. Переменные, созданные вне области, не будут распределены, что может повлиять на производительность. К распространённым объектам TensorFlow, создающим переменные, относятся модели, оптимизаторы и метрики. Такие объекты всегда следует инициализировать в области, а функции, которые могут создавать переменные отложенно (например, Model.call(), трассировка tf.function и т. д.), также следует вызывать внутри области. Переменные также могут создаваться при восстановлении из контрольной точки, если создание переменных происходит отложенно. Обратите внимание, что любая переменная, созданная внутри стратегии, сохраняет информацию о стратегии. Поэтому чтение и запись таких переменных вне strategy.scope также могут выполняться без проблем, и пользователю не требуется входить в область.
• Некоторые API стратегий (например, strategy.run и strategy.reduce), для использования которых необходимо находиться в области стратегии, входят в эту область автоматически. Это означает, что при использовании таких API вам не нужно входить в область самостоятельно.
• Когда tf.keras.Model создаётся внутри strategy.scope, объект модели сохраняет информацию об области. При последующем вызове методов высокоуровневой обучающей платформы, таких как 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 иногда необходимо выполнять внутри области, если при этом создаются переменные.