tf.distribute.experimental.TPUStrategy

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

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

tf.distribute.experimental.TPUStrategy(
    tpu_cluster_resolver=None, steps_per_run=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). В случаях, когда набор данных бесконечен, это разбиение можно выполнить, создав копии наборов данных, которые различаются только своим случайным seed. 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, session=None
)

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

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

Обратите внимание, что, вероятно, вам нужно использовать tf.distribute.Strategy.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

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

experimental_run(
    fn, input_iterator=None
)

Выполняет операции в fn на каждой реплике с входными данными из input_iterator.

УСТАРЕЛО: Этот метод недоступен в TF 2.x. Переключитесь на использование experimental_run_v2 вместо этого.

При включенном режиме выполнения операций, указанных в fn на каждой реплике. В противном случае создается граф для выполнения операций на каждой реплике.

Каждая реплика примет единственный, отличный вход из входных данных, предоставленных одним вызовом get_next на итераторе ввода.

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

experimental_run_v2

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

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

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

make_dataset_iterator

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

make_dataset_iterator(
    dataset
)

Создаёт итератор для входных данных, предоставленных через dataset.

УСТАРЕВШИЙ МЕТОД: Этот метод недоступен в TF 2.x.

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

Пользователь также может использовать make_input_fn_iterator, если хочет настроить, какие входные данные подаются на какую реплику/рабочий процесс и т. д.

make_input_fn_iterator

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

make_input_fn_iterator(
    input_fn, replication_mode=tf.distribute.InputReplicationMode.PER_WORKER
)

Возвращает итератор, разделенный по репликам, созданный из функции ввода.

УСТАРЕВШИЙ МЕТОД: Этот метод недоступен в TF 2.x.

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

def input_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)
with strategy.scope():
  iterator = strategy.make_input_fn_iterator(input_fn)
  replica_results = strategy.experimental_run(replica_fn, iterator)

Возвращаемый tf.data.Dataset input_fn должен иметь размер пакета для каждой реплики, который может быть вычислен с помощью input_context.get_per_replica_batch_size.

reduce

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

reduce(
    reduce_op, value, axis=None
)

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

update_config_proto

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

update_config_proto(
    config_proto
)

Возвращает копию config_proto, изменённую для использования с этой стратегией.

УСТАРЕВШИЙ МЕТОД: Этот метод недоступен в TF 2.x.

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