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 возвращаемого этим 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.

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

УСТАРЕЛО: Этот метод недоступен в TF 2.x. Переключитесь на использование run вместо него.

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

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

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

make_dataset_iterator

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

make_dataset_iterator(
    dataset
)

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

УСТАРЕЛО: Этот метод недоступен в TF 2.x.

Данные из предоставленного набора данных будут равномерно распределены по всем вычислительным репликам. Мы будем предполагать, что входной набор данных сгруппирован по глобальному размеру партии. С этим предположением мы постараемся разделить каждую партию по всем репликам (один или несколько рабочих узлов). Если эта попытка не увенчается успехом, будет выброшено исключение, и пользователю следует вместо этого использовать 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
)

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

УСТАРЕЛО: Этот метод недоступен в TF 2.x.

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

Что должно быть в области и что вне её?

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

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

update_config_proto

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

update_config_proto(
    config_proto
)

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

УСТАРЕВШИЙ: Этот метод недоступен в TF 2.x.

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