tf.distribute.experimental.TPUStrategy

Синхронное обучение на TPU и 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 на TPU можно использовать 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 объекта tf.distribute.DistributedDataset, возвращаемого этим API, позволяет запросить 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 на пакеты нового размера, равного глобальному размеру пакета, делённому на число синхронизированных реплик. Перебор выполняется с помощью цикла for в стиле 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.

По умолчанию этот метод добавляет преобразование предварительной выборки в конец предоставленного пользователем экземпляра 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>)
       

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

   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)
       

3. Задание значений с помощью 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)
       

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

   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. Для многоклиентской стратегии tf.distribute.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 и двумя 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 каждого рабочего узла.

Для редукции значений между репликами в tf.distribute предусмотрено несколько API:

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

Например, если глобальный размер пакета равен 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)

Если последний пакет неполный, нужно указать axis так, чтобы форма результата была одинаковой для всех реплик. Например, если размер последнего пакета равен 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() возвращает эту стратегию. За её пределами возвращается стратегия no-op по умолчанию.
• Вход в область также означает вход в «контекст между репликами». Объяснение контекстов между репликами и реплик см. в разделе 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 иногда необходимо выполнять внутри области действия, если при этом создаются переменные.