tf.compat.v2.distribute.experimental.TPUStrategy

Реализация стратегии распределения TPU.

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

tf.compat.v2.distribute.experimental.TPUStrategy(
    tpu_cluster_resolver=None, device_assignment=None
)

Методы

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
)

См. базовый класс.

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