tf.compat.v1.distribute.experimental.TPUStrategy

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

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

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

os.environ['TF_CONFIG'] = json.dumps({
  'cluster': {
      'worker': ["localhost:12345", "localhost:23456"],
      'ps': ["localhost:34567"]
  },
  'task': {'type': 'worker', 'index': 0}
})

# This implicitly uses TF_CONFIG for the cluster and current task info.
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()

...

if strategy.cluster_resolver.task_type == 'worker':
  # Perform something that's only applicable on workers. Since we set this
  # as a worker above, this block will run on this particular instance.
elif strategy.cluster_resolver.task_type == 'ps':
  # Perform something that's only applicable on parameter servers. Since we
  # set this as a worker above, this block will not run on this particular
  # instance.

Методы

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

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

experimental_local_results(
    value
)

Возвращает список всех локальных значений по реплике, содержащихся в value.

Примечание: Это возвращает только значения на рабочем процессе, инициированном этим клиентом. При использовании tf.distribute.Strategy, такого как tf.distribute.experimental.MultiWorkerMirroredStrategy, каждый рабочий процесс будет своим клиентом, и эта функция вернёт только значения, вычисленные на этом рабочем процессе.

experimental_make_numpy_dataset

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

experimental_make_numpy_dataset(
    numpy_input, session=None
)

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

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

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

Пример:

numpy_input = np.ones([10], dtype=np.float32)
dataset = strategy.experimental_make_numpy_dataset(numpy_input)
dist_dataset = strategy.experimental_distribute_dataset(dataset)

experimental_run

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

experimental_run(
    fn, input_iterator=None
)

Выполняет операторы в fn на каждой реплике, используя входные данные из input_iterator. (устарело)

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

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

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

make_dataset_iterator

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

make_dataset_iterator(
    dataset
)

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

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

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

make_input_fn_iterator

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

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

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

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

def input_fn(input_context):
  batch_size = input_context.get_per_replica_batch_size(global_batch_size)
  d = tf.data.Dataset.from_tensors([[1.]]).repeat().batch(batch_size)
  return d.shard(input_context.num_input_pipelines,
                 input_context.input_pipeline_id)
with strategy.scope():
  iterator = strategy.make_input_fn_iterator(input_fn)
  replica_results = strategy.experimental_run(replica_fn, iterator)

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

reduce

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

reduce(
    reduce_op, value, axis=None
)

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

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 графическими процессорами:

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, поскольку они позволяют настроить место назначения результата. Они также вызываются в контексте между репликами.

Каким должен быть ось?

Учитывая значение на реплику, возвращаемое 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
)

Запустить fn на каждой реплике с заданными аргументами.

Выполняет операции, указанные в fn на каждой реплике. Если args или kwargs имеют «значения на реплику», такие как те, которые производятся «распределённым Dataset», при выполнении fn на определённой реплике, оно будет выполнено с компонентом этих «значений на реплику», соответствующих этой реплике.

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

Все аргументы в args или kwargs должны быть либо вложенными тензорами, либо объектами на реплику, содержащими тензоры или составные тензоры.

Пользователи могут передать стратегические параметры в аргумент options. Пример для включения букетизации динамических форм в TPUStrategy.run:

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)

options = tf.distribute.RunOptions(
    experimental_bucketizing_dynamic_shape=True)

dataset = tf.data.Dataset.range(
   strategy.num_replicas_in_sync, output_type=dtypes.float32).batch(
       strategy.num_replicas_in_sync, drop_remainder=True)
input_iterator = iter(strategy.experimental_distribute_dataset(dataset))

@tf.function()
def step_fn(inputs):
 output = tf.reduce_sum(inputs)
 return output

strategy.run(step_fn, args=(next(input_iterator),), options=options)

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, чтобы явно распределить эти вычисления. См. пример в учебнике по пользовательской петле обучения custom training loop tutorial.

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

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

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

update_config_proto

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

update_config_proto(
    config_proto
)

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

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