tf.compat.v2.distribute.Strategy

Политика распределения состояния и вычислений по списку устройств.

tf.compat.v2.distribute.Strategy(
    extended
)

См. руководство для обзора и примеров.

Кратко:

• Для использования с Keras compile/fit, пожалуйста, ознакомьтесь.
• Вы можете передать наследника tf.distribute.Strategy в tf.estimator.RunConfig, чтобы указать, как tf.estimator.Estimator должен распределять свои вычисления. См. руководство.
• В противном случае, используйте tf.distribute.Strategy.scope, чтобы указать, что стратегия должна использоваться при построении и выполнении вашей модели. (Это помещает вас в «контекст между репликами» для этой стратегии, что означает, что стратегия контролирует такие вещи, как размещение переменных).

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

  • Начните с создания tf.data.Dataset обычным способом или используя tf.distribute.experimental_make_numpy_dataset для создания набора данных из массива numpy.

  • Используйте tf.distribute.Strategy.experimental_distribute_dataset для преобразования tf.data.Dataset в нечто, что генерирует значения «для каждой реплики». Если вы хотите вручную указать, как набор данных должен быть разделен между репликами, используйте tf.distribute.Strategy.experimental_distribute_datasets_from_function вместо этого.

  • Используйте tf.distribute.Strategy.experimental_run_v2 для выполнения функции один раз для каждой реплики, принимая значения, которые могут быть «для каждой реплики» (например, из распределенного набора данных) и возвращая значения «для каждой реплики». Эта функция выполняется в «контексте реплики», что означает, что каждая операция выполняется отдельно на каждой реплике.

  • Используйте метод (например, tf.distribute.Strategy.reduce), чтобы преобразовать полученные значения «для каждой реплики» в обычные Tensor.

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

with my_strategy.scope():
  @tf.function
  def distribute_train_epoch(dataset):
    def replica_fn(input):
      # process input and return result
      return result

    total_result = 0
    for x in dataset:
      per_replica_result = my_strategy.experimental_run_v2(replica_fn,
                                                           args=(x,))
      total_result += my_strategy.reduce(tf.distribute.ReduceOp.SUM,
                                         per_replica_result, axis=None)
    return total_result

  dist_dataset = my_strategy.experimental_distribute_dataset(dataset)
  for _ in range(EPOCHS):
    train_result = distribute_train_epoch(dist_dataset)

Это принимает обычный dataset и replica_fn и выполняет его в распределённом виде, используя определённую tf.distribute.Strategy с именем my_strategy выше. Любые переменные, созданные в replica_fn, создаются с политикой my_strategy, а библиотечные функции, вызываемые replica_fn, могут использовать API get_replica_context() для реализации распределённого поведения.

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

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

Примечание: tf.distribute.Strategy в настоящее время не поддерживает разнесённые переменные TensorFlow (где одна переменная разделена на несколько устройств).

Методы

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, и войдёт в свой "меж-репликационный контекст".