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.

Методы

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. 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.

По умолчанию этот метод добавляет преобразование prefetch в конец экземпляра tf.data.Dataset, предоставленного пользователем. Аргумент преобразования prefetch, который является 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-й размерности. Результат копируется на "текущее" устройство

• которое обычно является процессором (CPU) рабочего узла, на котором выполняется программа. Для tf.distribute.TPUStrategy это первый хост TPU. Для многоклиентских MultiWorkerMirroredStrategy, это процессор (CPU) каждого рабочего узла.

Эта 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 перед выводом.

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

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