tf.distribute.experimental.TPUStrategy

Синхронное обучение на TPUs и TPU-кластерах.

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

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

Для создания объекта TPUStrategy необходимо выполнить код инициализации, как показано ниже:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

При использовании стратегий распределения переменные, созданные в рамках области действия стратегии, будут дублироваться на всех репликах и могут поддерживаться синхронизированными с помощью алгоритмов all-reduce.

Для запуска программ TF2 на TPUs можно использовать API .compile и .fit в tf.keras со стратегией TPUStrategy или написать собственный настраиваемый цикл обучения, вызвав strategy.run напрямую. Обратите внимание, что TPUStrategy не поддерживает чистый eager-режим, поэтому убедитесь, что функция, переданная в strategy.run является tf.function, или strategy.run вызывается внутри tf.function, если включено поведение eager.

Методы

experimental_assign_to_logical_device

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

experimental_assign_to_logical_device(
    tensor, logical_device_id
)

Добавляет аннотацию, что tensor будет назначен логическому устройству.

Примечание: Этот API в настоящее время поддерживается только в TPUStrategy. Это добавляет аннотацию к tensor, указывающую, что операции над tensor будут вызываться на логическом устройстве с ID logical_device_id. При использовании моделирования параллельности по умолчанию все операции размещаются на логическом устройстве с индексом ноль.

# Initializing TPU system with 2 logical devices and 4 replicas.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 1, 1, 2],
    num_replicas=4)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)
iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  output = tf.add(inputs, inputs)

  # Add operation will be executed on logical device 0.
  output = strategy.experimental_assign_to_logical_device(output, 0)
  return output

strategy.run(step_fn, args=(next(iterator),))

experimental_distribute_dataset

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

experimental_distribute_dataset(
    dataset, options=None
)

Создает tf.distribute.DistributedDataset из tf.data.Dataset.

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

Вот пример:

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 `tf.distribute.DistributedDataset`
for x in dist_dataset:
  # process dataset elements
  strategy.run(replica_fn, args=(x,))

В приведенном выше фрагменте кода tf.distribute.DistributedDataset dist_dataset группируется по GLOBAL_BATCH_SIZE, и мы перебираем его с помощью for x in dist_dataset. x — tf.distribute.DistributedValues, содержащий данные для всех реплик, которые агрегируются в пакет из GLOBAL_BATCH_SIZE. tf.distribute.Strategy.run позаботится о подаче правильных данных для каждой реплики в x в соответствующую replica_fn , выполняемую на каждой реплике.

Что происходит за кулисами этого метода, когда мы говорим, что экземпляр tf.data.Dataset - dataset - распределяется? Это зависит от того, как вы настраиваете tf.data.experimental.AutoShardPolicy через tf.data.experimental.DistributeOptions. По умолчанию он установлен в tf.data.experimental.AutoShardPolicy.AUTO. В многоузловой системе мы сначала попытаемся распределить dataset путем обнаружения, создается ли dataset из наборов данных для чтения (например, tf.data.TFRecordDataset, tf.data.TextLineDataset и т. д.) и, если да, попытаемся разбить входные файлы. Обратите внимание, что должен быть как минимум один входной файл на каждый узел. Если у вас меньше одного входного файла на узел, мы рекомендуем отключить фрагментацию наборов данных между узлами, установив tf.data.experimental.DistributeOptions.auto_shard_policy в значение tf.data.experimental.AutoShardPolicy.OFF.

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

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

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

Вы также можете использовать свойство element_spec экземпляра tf.distribute.DistributedDataset, возвращенного этим API, чтобы запросить tf.TypeSpec элементов, возвращаемых итератором. Это можно использовать для установки свойства input_signature tf.function.

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)

@tf.function(input_signature=[dist_dataset.element_spec])
def train_step(inputs):
  # train model with inputs
  return

# Iterate over the `tf.distribute.DistributedDataset`
for x in dist_dataset:
  # process dataset elements
  strategy.run(train_step, args=(x,))

Примечание: Порядок обработки данных рабочими узлами при использовании tf.distribute.Strategy.experimental_distribute_dataset или tf.distribute.Strategy.experimental_distribute_datasets_from_function не гарантируется. Это обычно требуется, если вы используете tf.distribute для масштабирования прогнозирования. Однако вы можете вставить индекс для каждого элемента в пакет и упорядочить выходные данные соответственно. Обратитесь к этому фрагменту для примера того, как упорядочить выходные данные.

experimental_distribute_datasets_from_function

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

experimental_distribute_datasets_from_function(
    dataset_fn, options=None
)

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

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

Этот метод можно использовать для нескольких целей. Например, где 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.

global_batch_size = 8
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)

strategy = tf.distribute.MirroredStrategy()
ds = strategy.experimental_distribute_datasets_from_function(dataset_fn)

def train(ds):
  @tf.function(input_signature=[ds.element_spec])
  def step_fn(inputs):
    # train the model with inputs
    return inputs

... для пакета в ds: ... replica_results = strategy.run(replica_fn, args=(batch,))

train(ds)

Примечание: Порядок, в котором данные обрабатываются рабочими процессами при использовании tf.distribute.Strategy.experimental_distribute_dataset или tf.distribute.Strategy.experimental_distribute_datasets_from_function, не гарантируется. Это обычно требуется, если вы используете tf.distribute для масштабирования прогнозирования. Однако вы можете вставить индекс для каждого элемента в пакете и упорядочить выводы соответственно. Обратитесь к этому фрагменту для примера того, как упорядочить выводы.

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()
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>,)

1. Распределение значений в массиве на основе replica_id:

strategy = tf.distribute.MirroredStrategy()
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,)

1. Указание значений с помощью num_replicas_in_sync:

strategy = tf.distribute.MirroredStrategy()
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
(1,)

1. Размещение значений на устройствах и распределение:

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, каждый работник будет своим клиентом, и эта функция вернет только значения, вычисленные на этом работнике.

experimental_make_numpy_dataset

Просмотр исходного кода

experimental_make_numpy_dataset(
    numpy_input
)

Создает tf.data.Dataset из массива NumPy. (устарело)

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

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

Пример:

strategy = tf.distribute.MirroredStrategy()
numpy_input = np.ones([10], dtype=np.float32)
dataset = strategy.experimental_make_numpy_dataset(numpy_input)
dataset
<TensorSliceDataset shapes: (), types: tf.float32>
dataset = dataset.batch(2)
dist_dataset = strategy.experimental_distribute_dataset(dataset)

experimental_replicate_to_logical_devices

Просмотр исходного кода

experimental_replicate_to_logical_devices(
    tensor
)

Добавляет аннотацию, что tensor будет дублироваться на все логические устройства.

Примечание: Этот API поддерживается только в TPUStrategy на данный момент. Это добавляет аннотацию к тензору tensor, указывающую, что операции с tensor будут вызваны на всех логических устройствах.

# Initializing TPU system with 2 logical devices and 4 replicas.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 1, 1, 2],
    num_replicas=4)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)

iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  images, labels = inputs
  images = strategy.experimental_split_to_logical_devices(
    inputs, [1, 2, 4, 1])

  # model() function will be executed on 8 logical devices with `inputs`
  # split 2 * 4  ways.
  output = model(inputs)

  # For loss calculation, all logical devices share the same logits
  # and labels.
  labels = strategy.experimental_replicate_to_logical_devices(labels)
  output = strategy.experimental_replicate_to_logical_devices(output)
  loss = loss_fn(labels, output)

  return loss

strategy.run(step_fn, args=(next(iterator),))

Args: tensor: Входной тензор для аннотации.

experimental_split_to_logical_devices

Просмотр исходного кода

experimental_split_to_logical_devices(
    tensor, partition_dimensions
)

Добавляет аннотацию, что tensor будет разделен по логическим устройствам.

Примечание: Этот API поддерживается только в TPUStrategy на данный момент. Это добавляет аннотацию к тензору tensor, указывающую, что операции с tensor будут разделены между несколькими логическими устройствами. Тензор tensor будет разделен по измерениям, указанным в partition_dimensions. Измерения tensor должны быть кратны соответствующему значению в partition_dimensions.

Например, для системы с 8 логическими устройствами, если tensor - это тензор изображения со формой (размер_пакета, ширина, высота, канал) и partition_dimensions - это [1, 2, 4, 1], то tensor будет разделен 2 способами по ширине и 4 способами по высоте, а значения разбиения тензора будут переданы на 8 логических устройств.

# Initializing TPU system with 8 logical devices and 1 replica.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 2, 2, 2],
    num_replicas=1)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)

iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  inputs = strategy.experimental_split_to_logical_devices(
    inputs, [1, 2, 4, 1])

  # model() function will be executed on 8 logical devices with `inputs`
  # split 2 * 4  ways.
  output = model(inputs)
  return output

strategy.run(step_fn, args=(next(iterator),))

Args: tensor: Входной тензор для аннотации. partition_dimensions: Не вложенный список целых чисел размером, равным рангу tensor , определяющий, как будет разбит tensor . Произведение всех элементов в partition_dimensions должно быть равно общему числу логических устройств на реплику.

reduce

Просмотр исходного кода

reduce(
    reduce_op, value, axis
)

Сведение value по репликам.

Учитывая значение на реплику, возвращаемое run, скажем, потерю на пример, пакет будет разделен между всеми репликами. Эта функция позволяет агрегировать по репликам и необязательно также по элементам пакета. Например, если у вас есть глобальный размер пакета 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.

data-text="run" id="run" tabindex="0">run

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

run(
    fn, args=(), kwargs=None, options=None
)

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

data-text="scope" id="scope" tabindex="0">scope

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

scope()

Контекстный менеджер для установления стратегии как текущей и распределения переменных.

Этот метод возвращает контекстный менеджер и используется следующим образом:

strategy = tf.distribute.MirroredStrategy()
# 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>
}
# 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 для явного распределения этого вычисления. См. пример в руководстве по созданию пользовательской обучающей петли .

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

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

• Все, что создаёт переменные, которые должны быть распределёнными переменными, должно находиться в strategy.scope. Это может быть сделано либо путем непосредственного помещения в область, либо с помощью другого API, например, strategy.run или model.fit, чтобы он попал туда за вас. Любая переменная, созданная вне области действия, не будет распределена и может иметь последствия для производительности. Общие вещи, создающие переменные в TF: модели, оптимизаторы, метрики. Они всегда должны создаваться в области действия. Ещё одним источником создания переменных может быть восстановление контрольной точки - когда переменные создаются лениво. Обратите внимание, что любая переменная, созданная внутри стратегии, сохраняет информацию о стратегии. Поэтому чтение и запись в эти переменные вне 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 иногда может потребоваться находиться внутри области действия, если оно создаёт переменные.