tf.compat.v2.distribute.experimental.MultiWorkerMirroredStrategy

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

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

tf.compat.v2.distribute.experimental.MultiWorkerMirroredStrategy(
    communication=tf.distribute.experimental.CollectiveCommunication.AUTO
)

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

Она использует реализацию multi-worker all-reduce CollectiveOps для синхронизации переменных. Коллективная операция — это одиночная операция в графе TensorFlow, которая может автоматически выбирать алгоритм all-reduce в runtime TensorFlow в зависимости от оборудования, топологии сети и размеров тензоров.

По умолчанию она использует все локальные графические процессоры или ЦП для обучения на одном рабочем узле.

Когда переменная среды 'TF_CONFIG' установлена, она анализирует cluster_spec, task_type и task_id из 'TF_CONFIG' и преобразует их в стратегию для нескольких рабочих узлов, которая дублирует модели на графических процессорах всех машин в кластере. В текущей реализации она использует все графические процессоры в кластере и предполагает, что все рабочие узлы имеют одинаковое количество графических процессоров.

Она поддерживает как режим eager, так и режим графа. Однако для режима eager ей необходимо настроить контекст eager в своём конструкторе, и поэтому все операции в режиме eager должны выполняться после создания объекта стратегии.

Краткое описание методов

experimental_distribute_dataset

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

experimental_distribute_dataset(
    dataset
)

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

Возвращаемый распределённый набор данных может быть проитерирован аналогично обычным наборам данных. ПРИМЕЧАНИЕ: В настоящее время пользователь не может добавить больше преобразований к распределённому набору данных.

Следующий пример:

strategy = tf.distribute.MirroredStrategy()

# Create a dataset
dataset = dataset_ops.Dataset.TFRecordDataset([
  "/a/1.tfr", "/a/2.tfr", "/a/3.tfr", "/a/4.tfr"])

# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(dataset)
# Iterate over the distributed dataset
for x in dist_dataset:
  # process dataset elements
  strategy.experimental_run_v2(train_step, args=(x,))

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

В многоузловой среде мы сначала попытаемся распределить набор данных, попытавшись определить, создаётся ли он из ReaderDatasets (например, TFRecordDataset, TextLineDataset и т. д.), и если да, то попытаемся разбить входные файлы. Обратите внимание, что должно быть по крайней мере один входной файл на рабочий узел. Если у вас меньше одного входного файла на рабочий узел, мы рекомендуем отключить распределение вашего набора данных, используя описанный ниже метод.

Если эта попытка окажется неудачной (например, набор данных создаётся из Dataset.range), мы разделим набор данных равномерно в конце, добавив операцию .shard в конец потока обработки. Это приведёт к запуску всей предварительной обработки всего данных на каждом рабочем узле, и каждый рабочий узел будет выполнять избыточную работу. Мы выведем предупреждение, если этот метод разбиения будет выбран. В этом случае рассмотрите использование experimental_distribute_datasets_from_function вместо него.

Вы можете отключить фрагментацию наборов данных между рабочими узлами, используя параметр auto_shard в tf.data.experimental.DistributeOptions.

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

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

experimental_distribute_datasets_from_function

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

experimental_distribute_datasets_from_function(
    dataset_fn
)

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

dataset_fn будет вызываться один раз для каждого рабочего узла в стратегии. Каждая реплика на этом рабочем узле будет извлекать одну партию ввода из локального Dataset (то есть, если на рабочем узле есть две реплики, две партии будут извлекаться из Dataset на каждом шаге).

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

Функция dataset_fn должна принимать экземпляр tf.distribute.InputContext, где можно получить информацию о пакетировании и репликации ввода:

def dataset_fn(input_context):
  batch_size = input_context.get_per_replica_batch_size(global_batch_size)
  d = tf.data.Dataset.from_tensors([[1.]]).repeat().batch(batch_size)
  return d.shard(
      input_context.num_input_pipelines, input_context.input_pipeline_id)

inputs = strategy.experimental_distribute_datasets_from_function(dataset_fn)

for batch in inputs:
  replica_results = strategy.experimental_run_v2(replica_fn, args=(batch,))

experimental_local_results

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

experimental_local_results(
    value
)

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

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

experimental_make_numpy_dataset

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

experimental_make_numpy_dataset(
    numpy_input
)

Создаёт tf.data.Dataset для входных данных, предоставленных через массив NumPy.

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

Обратите внимание, что вам, вероятно, потребуется использовать experimental_distribute_dataset с возвращаемым набором данных для дальнейшего распределения его с помощью стратегии.

Пример:

numpy_input = np.ones([10], dtype=np.float32)
dataset = strategy.experimental_make_numpy_dataset(numpy_input)
dist_dataset = strategy.experimental_distribute_dataset(dataset)

experimental_run_v2

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

experimental_run_v2(
    fn, args=(), kwargs=None
)

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

Выполняет операции, указанные в fn, на каждой реплике. Если args или kwargs имеют значения «на реплику», такие как те, которые генерирует «распределённый Dataset», когда fn выполняется на конкретной реплике, она будет выполняться с компонентом этих значений «на реплику», соответствующим этой реплике.

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

Все аргументы в args или kwargs должны быть либо вложенными тензорами, либо объектами «на реплику», содержащими тензоры или составные тензоры.

reduce

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

reduce(
    reduce_op, value, axis
)

Свести value по репликам.

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

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

scope

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

scope()

Возвращает менеджер контекста, выбирающий эту стратегию как текущую.

Внутри блока кода with strategy.scope():, эта нить будет использовать создатель переменных, установленный strategy, и войдёт в свой «межреплицированный контекст».

В MultiWorkerMirroredStrategy, все переменные, созданные внутри `strategy.scope()`, будут дублироваться на всех репликах каждого рабочего узла. Кроме того, он также устанавливает область контекста устройства по умолчанию, поэтому операции без указанных устройств будут выполняться на правильном рабочем узле.