tf.distribute.TPUStrategy

Синхронное обучение на TPU и TPU Pod.

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

tf.distribute.TPUStrategy(
    tpu_cluster_resolver=None,
    experimental_device_assignment=None,
    experimental_spmd_xla_partitioning=False
)

Используется в блокнотах

Чтобы создать объект 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.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-режима. Подробнее см. на странице https://www.tensorflow.org/guide/tpu.

API distribute_datasets_from_function и experimental_distribute_dataset можно использовать для распределения набора данных между рабочими узлами TPU при написании собственного цикла обучения. Если вы используете методы fit и compile, доступные в tf.keras.Model, распределением займётся Keras.

Пример написания собственного цикла обучения на TPU:

with strategy.scope():
  model = tf.keras.Sequential([
    tf.keras.layers.Dense(2, input_shape=(5,)),
  ])
  optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)

def dataset_fn(ctx):
  x = np.random.random((2, 5)).astype(np.float32)
  y = np.random.randint(2, size=(2, 1))
  dataset = tf.data.Dataset.from_tensor_slices((x, y))
  return dataset.repeat().batch(1, drop_remainder=True)
dist_dataset = strategy.distribute_datasets_from_function(
    dataset_fn)
iterator = iter(dist_dataset)

@tf.function()
def train_step(iterator):

  def step_fn(inputs):
    features, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(features, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits)

    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

  strategy.run(step_fn, args=(next(iterator),))

train_step(iterator)

Для более сложных сценариев, например параллелизма моделей, при создании TPUStrategy можно задать аргумент experimental_device_assignment, чтобы указать число реплик и логических устройств. Ниже приведён пример инициализации системы TPU с двумя логическими устройствами и одной репликой.

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 1, 1, 2],
    num_replicas=1)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)

Затем можно выполнить операцию tf.add только на логическом устройстве 0.

@tf.function()
def step_fn(inputs):
  features, _ = inputs
  output = tf.add(features, features)

  # Add operation will be executed on logical device 0.
  output = strategy.experimental_assign_to_logical_device(output, 0)
  return output
dist_dataset = strategy.distribute_datasets_from_function(
    dataset_fn)
iterator = iter(dist_dataset)
strategy.run(step_fn, args=(next(iterator),))

experimental_spmd_xla_partitioning включает экспериментальную функцию XLA SPMD для параллелизма моделей. Этот флаг может сократить время компиляции и требования к HBM. В этом режиме каждый входной тензор должен быть либо разделён на части (с помощью strategy.experimental_split_to_logical_devices), либо полностью реплицирован (с помощью strategy.experimental_replicate_to_logical_devices) на все логические устройства. Вызов strategy.experimental_assign_to_logical_device в этом режиме приведёт к ValueError.

Методы

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_assign_to_logical_device

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

experimental_assign_to_logical_device(
    tensor, logical_device_id
)

Добавляет аннотацию о назначении tensor логическому устройству.

К tensor добавляется аннотация, указывающая, что операции с tensor будут выполняться на логическом устройстве с идентификатором logical_device_id. При использовании параллелизма моделей по умолчанию все операции размещаются на логическом устройстве с индексом ноль.

# Initializing TPU system with 2 logical devices and 4 replicas.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 1, 1, 2],
    num_replicas=4)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)
iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  output = tf.add(inputs, inputs)

  # Add operation will be executed on logical device 0.
  output = strategy.experimental_assign_to_logical_device(output, 0)
  return output

strategy.run(step_fn, args=(next(iterator),))

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

experimental_replicate_to_logical_devices

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

experimental_replicate_to_logical_devices(
    tensor
)

Добавляет аннотацию о репликации tensor на всех логических устройствах.

К тензору tensor добавляется аннотация, указывающая, что операции с tensor будут выполняться на всех логических устройствах.

# Initializing TPU system with 2 logical devices and 4 replicas.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 1, 1, 2],
    num_replicas=4)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)

iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  images, labels = inputs
  images = strategy.experimental_split_to_logical_devices(
    inputs, [1, 2, 4, 1])

  # model() function will be executed on 8 logical devices with `inputs`
  # split 2 * 4  ways.
  output = model(inputs)

  # For loss calculation, all logical devices share the same logits
  # and labels.
  labels = strategy.experimental_replicate_to_logical_devices(labels)
  output = strategy.experimental_replicate_to_logical_devices(output)
  loss = loss_fn(labels, output)

  return loss

strategy.run(step_fn, args=(next(iterator),))

Аргументы: tensor: входной тензор, которому нужно назначить аннотацию.

experimental_split_to_logical_devices

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

experimental_split_to_logical_devices(
    tensor, partition_dimensions
)

Добавляет аннотацию о разделении tensor между логическими устройствами.

К тензору tensor добавляется аннотация, указывающая, что операции с tensor будут распределены между несколькими логическими устройствами. Тензор tensor будет разделён по измерениям, заданным в partition_dimensions. Размеры tensor должны делиться на соответствующие значения в partition_dimensions.

Например, для системы с 8 логическими устройствами, если tensor — это тензор изображения с формой (batch_size, width, height, channel), а partition_dimensions — [1, 2, 4, 1], то tensor будет разделён на 2 части по измерению width и на 4 части по измерению height, а значения разделённого тензора будут переданы 8 логическим устройствам.

# Initializing TPU system with 8 logical devices and 1 replica.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 2, 2, 2],
    num_replicas=1)
# Construct the TPUStrategy. Since we are going to split the image across
# logical devices, here we set `experimental_spmd_xla_partitioning=True`
# so that the partitioning can be compiled in SPMD mode, which usually
# results in faster compilation and smaller HBM requirement if the size of
# input and activation tensors are much bigger than that of the model
# parameters. Note that this flag is suggested but not a hard requirement
# for `experimental_split_to_logical_devices`.
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment,
    experimental_spmd_xla_partitioning=True)

iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  inputs = strategy.experimental_split_to_logical_devices(
    inputs, [1, 2, 4, 1])

  # model() function will be executed on 8 logical devices with `inputs`
  # split 2 * 4  ways.
  output = model(inputs)
  return output

strategy.run(step_fn, args=(next(iterator),))

Аргументы: tensor: входной тензор, которому нужно назначить аннотацию. partition_dimensions: плоский список целых чисел, размер которого равен рангу tensor; список задаёт способ разделения tensor. Произведение всех элементов partition_dimensions должно равняться общему числу логических устройств на реплику.

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, а реплик две, значения для примеров [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, на каждой реплике TPU.

Выполняет операции, заданные fn, на каждой реплике. Если args или kwargs содержат tf.distribute.DistributedValues, например созданные объектом tf.distribute.DistributedDataset с помощью tf.distribute.Strategy.experimental_distribute_dataset или tf.distribute.Strategy.distribute_datasets_from_function, то при выполнении fn на определенной реплике будут использоваться компоненты tf.distribute.DistributedValues, соответствующие этой реплике.

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

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

Пример использования:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
@tf.function
def run():
  def value_fn(value_context):
    return value_context.num_replicas_in_sync
  distributed_values = (
      strategy.experimental_distribute_values_from_function(value_fn))
  def replica_fn(input):
    return input * 2
  return strategy.run(replica_fn, args=(distributed_values,))
result = run()

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 или keras.Model.fit, который автоматически выполнит вход в область. Переменные, созданные вне области, не будут распределенными, что может повлиять на производительность. К распространенным объектам TensorFlow, создающим переменные, относятся модели, оптимизаторы и метрики. Такие объекты следует инициализировать в области, а любые функции, которые могут создавать переменные отложенно (например, 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 иногда нужно выполнять внутри области, если при этом создаются переменные.