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 в новую пачку размером, равным глобальному размеру пачки, деленному на количество реплик в синхронизации. Мы итерируем по нему с помощью питоновского цикла for. x — это tf.distribute.DistributedValues, содержащий данные для всех реплик, и каждая реплика получает данные нового размера пачки. tf.distribute.Strategy.run позаботится о подаче правильных данных для каждой реплики в x нужным replica_fn, выполняемым на каждой реплике.

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

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

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

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

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

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

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

experimental_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 для явного распределения вычислений. Пример смотрите в учебнике по пользовательскому циклу обучения https://www.tensorflow.org/tutorials/distribute/custom_training.

Что должно находиться в области действия, а что — вне её?

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

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

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