tf.distribute.experimental.TPUStrategy

Синхронное обучение на TPUs и TPU Pod.

Наследуется от: 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-поведение.

Методы

distribute_datasets_from_function

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

distribute_datasets_from_function(
    dataset_fn, options=None
)

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

Аргумент dataset_fn, передаваемый пользователем, представляет собой функцию ввода, имеющую аргумент tf.distribute.InputContext и возвращающую экземпляр tf.data.Dataset. Ожидается, что возвращаемый из dataset_fn набор данных уже сгруппирован по размеру пакета на реплику (т. е. глобальный размер пакета, деленный на количество реплик в синхронизации) и поделен. tf.distribute.Strategy.distribute_datasets_from_function не группирует и не разбивает экземпляр tf.data.Dataset, возвращаемый функцией ввода. dataset_fn будет вызываться на процессоре CPU каждого из рабочих узлов, и каждый из них генерирует набор данных, в котором каждая реплика на этом узле будет декьюить одну партию входных данных (т. е. если узел имеет две реплики, с набора данных Dataset будут декьюированы две партии на каждом шаге).

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

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

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

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

Примечание: Трансформации состояний наборов данных в настоящее время не поддерживаются с tf.distribute.experimental_distribute_dataset или tf.distribute.distribute_datasets_from_function. Любые операции состояния, которые могут иметь набор данных, в настоящее время игнорируются. Например, если ваш набор данных имеет map_fn, который использует tf.random.uniform для поворота изображения, то у вас есть граф набора данных, зависящий от состояния (т. е. случайного начального значения) на локальной машине, где выполняется процесс Python.

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

experimental_distribute_dataset

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

experimental_distribute_dataset(
    dataset, options=None
)

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

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

Вот пример:

global_batch_size = 2
# Passing the devices is optional.
strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
# Create a dataset
dataset = tf.data.Dataset.range(4).batch(global_batch_size)
# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(dataset)
@tf.function
def replica_fn(input):
  return input*2
result = []
# Iterate over the `tf.distribute.DistributedDataset`
for x in dist_dataset:
  # process dataset elements
  result.append(strategy.run(replica_fn, args=(x,)))
print(result)
[PerReplica:{
  0: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([0])>,
  1: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([2])>
}, PerReplica:{
  0: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([4])>,
  1: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([6])>
}]

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

В приведённом выше фрагменте кода dataset сгруппирован по global_batch_size, и вызов experimental_distribute_dataset на нём перегруппирует dataset в новый размер пакета, равный глобальному размеру пакета, делённому на количество реплик в синхронизации. Мы перебираем его с помощью цикла Python for. x — это tf.distribute.DistributedValues, содержащий данные для всех реплик, и каждая реплика получает данные нового размера пакета. tf.distribute.Strategy.run позаботится о предоставлении правильных данных на реплику в x нужной replica_fn на каждой реплике.

Разделение включает автоматическое разделение по нескольким рабочим узлам и внутри каждого рабочего узла. Во-первых, в распределённом обучении на нескольких рабочих узлах (т. е. когда вы используете tf.distribute.experimental.MultiWorkerMirroredStrategy или tf.distribute.TPUStrategy) автоматическое разделение набора данных по рабочим узлам означает, что каждому рабочему узлу назначается подмножество всего набора данных (если установлен правильный tf.data.experimental.AutoShardPolicy). Это гарантирует, что на каждом шаге глобальный размер пакета элементов непересекающихся наборов данных будет обрабатываться каждым рабочим узлом. Автоматическое разделение имеет несколько разных вариантов, которые можно указать с помощью tf.data.experimental.DistributeOptions. Далее, разделение внутри каждого рабочего узла означает, что метод разделит данные между всеми устройствами рабочего узла (если их несколько). Это произойдёт независимо от автоматического разделения по нескольким рабочим узлам.

Примечание: для автоматического разделения по нескольким рабочим узлам режим по умолчанию — tf.data.experimental.AutoShardPolicy.AUTO. Этот режим попытается разделить набор данных по файлам, если набор данных создаётся из наборов данных для чтения (например, tf.data.TFRecordDataset, tf.data.TextLineDataset и т. д.), или разделить набор данных по данным, где каждый из рабочих узлов прочитает весь набор данных и обработает только назначенное ему разбиение. Однако если у вас меньше одного входного файла на рабочий узел, мы рекомендуем отключить автоматическое разделение наборов данных по рабочим узлам, установив tf.data.experimental.DistributeOptions.auto_shard_policy в tf.data.experimental.AutoShardPolicy.OFF.

По умолчанию этот метод добавляет преобразование предварительной выборки в конце предоставленного пользователем экземпляра tf.data.Dataset. Аргументом преобразования предварительной выборки, который является buffer_size, является количество реплик в синхронизации.

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

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

Примечание: Преобразования наборов данных со состоянием в настоящее время не поддерживаются с tf.distribute.experimental_distribute_dataset или tf.distribute.distribute_datasets_from_function. Любые операторы со стетовым состоянием, которые могут быть у набора данных, в настоящее время игнорируются. Например, если у вашего набора данных есть map_fn, который использует tf.random.uniform для вращения изображения, то у вас есть граф набора данных, который зависит от состояния (т.е. случайного зерна) на локальной машине, где выполняется процесс Python.

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

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(["GPU:0", "GPU:1"])
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>,
 <tf.Tensor: shape=(), dtype=float32, numpy=1.0>)

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

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
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, 2.0)

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

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
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
(2, 2)

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

gather

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

gather(
    value, axis
)

Сборка value по репликам вдоль axis на текущее устройство.

Учитывая tf.distribute.DistributedValues или tf.Tensor-подобный объект value, этот API собирает и конкатенирует value по репликам вдоль axis-й размерности. Результат копируется на «текущее» устройство, которое обычно является процессором рабочего процесса, на котором выполняется программа. Для tf.distribute.TPUStrategy это первый хост TPU. Для многоклиентского tf.distribute.MultiWorkerMirroredStrategy это процессор каждого рабочего процесса.

Этот API может быть вызван только в контексте межрепликации. Для аналога в контексте реплики см. tf.distribute.ReplicaContext.all_gather.

Примечание: Для всех стратегий, кроме tf.distribute.TPUStrategy, вход value на разных репликах должен иметь одинаковый ранг, и их формы должны быть одинаковыми во всех размерностях, кроме axis-й размерности. Другими словами, их формы не могут отличаться в измерении d, где d не равно аргументу axis. Например, задан tf.distribute.DistributedValues с компонентами тензоров формы (1, 2, 3) и (1, 3, 3) на двух репликах, вы можете вызвать gather(..., axis=1, ...) на нем, но не gather(..., axis=0, ...) или gather(..., axis=2, ...). Однако для tf.distribute.TPUStrategy.gather все тензоры должны иметь ровно тот же ранг и ту же форму.

Примечание: Учитывая tf.distribute.DistributedValues value, его компоненты тензоров должны иметь ненулевой ранг. В противном случае рассмотрите использование tf.expand_dims перед их сбором.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
# A DistributedValues with component tensor of shape (2, 1) on each replica
distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(tf.constant([[1], [2]])))
@tf.function
def run():
  return strategy.gather(distributed_values, axis=0)
run()
<tf.Tensor: shape=(4, 1), dtype=int32, numpy=
array([[1],
       [2],
       [1],
       [2]], dtype=int32)>

Рассмотрим следующий пример для других комбинаций:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1", "GPU:2", "GPU:3"])
single_tensor = tf.reshape(tf.range(6), shape=(1,2,3))
distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(single_tensor))
@tf.function
def run(axis):
  return strategy.gather(distributed_values, axis=axis)
axis=0
run(axis)
<tf.Tensor: shape=(4, 2, 3), dtype=int32, numpy=
array([[[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]]], dtype=int32)>
axis=1
run(axis)
<tf.Tensor: shape=(1, 8, 3), dtype=int32, numpy=
array([[[0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5]]], dtype=int32)>
axis=2
run(axis)
<tf.Tensor: shape=(1, 2, 12), dtype=int32, numpy=
array([[[0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2],
        [3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5]]], dtype=int32)>

reduce

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

reduce(
    reduce_op, value, axis
)

Сведение value по репликам и возвращение результата на текущем устройстве.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
def step_fn():
  i = tf.distribute.get_replica_context().replica_id_in_sync_group
  return tf.identity(i)

per_replica_result = strategy.run(step_fn)
total = strategy.reduce("SUM", per_replica_result, axis=None)
total
<tf.Tensor: shape=(), dtype=int32, numpy=1>

Чтобы увидеть, как это будет выглядеть с несколькими репликами, рассмотрите тот же пример с MirroredStrategy с 2 GPU:

strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
def step_fn():
  i = tf.distribute.get_replica_context().replica_id_in_sync_group
  return tf.identity(i)

per_replica_result = strategy.run(step_fn)
# Check devices on which per replica result is:
strategy.experimental_local_results(per_replica_result)[0].device
# /job:localhost/replica:0/task:0/device:GPU:0
strategy.experimental_local_results(per_replica_result)[1].device
# /job:localhost/replica:0/task:0/device:GPU:1

total = strategy.reduce("SUM", per_replica_result, axis=None)
# Check device on which reduced result is:
total.device
# /job:localhost/replica:0/task:0/device:CPU:0

Этот API обычно используется для агрегирования результатов, возвращаемых из разных реплик, для отчетов и т. д. Например, потерю, вычисленную из разных реплик, можно усреднить с помощью этого API перед печатью.

Примечание: Результат копируется на «текущее» устройство — обычно на процессор рабочего процесса, на котором выполняется программа. Для TPUStrategy, это первый хост TPU. Для многоклиентского MultiWorkerMirroredStrategy, это процессор каждого рабочего процесса.

Существует ряд различных API tf.distribute для сведения значений по репликам:

• tf.distribute.ReplicaContext.all_reduce: Это отличается от Strategy.reduce тем, что предназначено для контекста реплики и не копирует результаты на устройство хоста. all_reduce обычно используется для вычислений внутри шага обучения, таких как градиенты.
• tf.distribute.StrategyExtended.reduce_to и tf.distribute.StrategyExtended.batch_reduce_to: Эти API являются более продвинутыми версиями Strategy.reduce, поскольку они позволяют настраивать место назначения результата. Они также вызываются в межрепликатоном контексте.

Что должно быть значением axis?

При заданном значении на реплику, возвращаемом run, например, потери на пример, пакет будет разделен между всеми репликами. Эта функция позволяет агрегировать по репликам и, по желанию, также по элементам пакета, указав параметр оси соответственно.

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

strategy.reduce("sum", per_replica_result, axis=None)

Иногда вам нужно будет агрегировать как по глобальному пакету, так и по всем репликам. Этого поведения можно добиться, указав размерность пакета как axis, обычно axis=0. В этом случае он вернет скаляр 0+1+2+3+4+5+6+7.

strategy.reduce("sum", per_replica_result, axis=0)

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

run

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

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

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

scope

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

scope()

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

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

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
# 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>,
  1: <tf.Variable 'Variable/replica_1: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 или keras.Model.fit, чтобы он входил в нее автоматически. Любая переменная, созданная за пределами области, не будет распределена и может иметь последствия для производительности. Некоторые общие объекты, которые создают переменные в TF, это модели, оптимизаторы, метрики. Такие объекты всегда должны инициализироваться в области, а любые функции, которые могут лениво создавать переменные (например, Model.__call__(), отслеживание tf.function и т. д.), аналогичным образом должны вызываться в области. Еще одним источником создания переменных может быть восстановление контрольной точки - когда переменные создаются лениво. Обратите внимание, что любая переменная, созданная внутри стратегии, фиксирует информацию о стратегии. Поэтому чтение и запись в эти переменные за пределами области strategy.scope также могут работать беспрепятственно без необходимости ввода области пользователем.
• Некоторые API стратегии (например, strategy.run и strategy.reduce) требующие размещения в области стратегии, входят в область автоматически, что означает, что при использовании этих API вам не нужно входить в область явно.
• Когда tf.keras.Model создается внутри strategy.scope, объект Model фиксирует информацию о области. При вызове высокоуровневых методов обучения, таких как 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 иногда может потребоваться внутри области, если оно создает переменные.