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

Если последний пакет неполный, необходимо указать axis, чтобы форма результата была одинаковой для всех реплик. Например, если в последнем пакете 6 элементов и он разделён на [0, 1, 2, 3] и [4, 5], возникнет несоответствие форм, если не указать axis=0. Если указать tf.distribute.ReduceOp.MEAN, то использование axis=0 обеспечит правильный знаменатель, равный 6. Сравните это с вычислением reduce_mean, которое даёт скалярное значение для каждой реплики, и последующим усреднением этих значений с помощью данной функции: при таком подходе некоторые значения получат вес 1/8, а другие — 1/4.

run

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

run(
    fn, args=(), kwargs=None, options=None
)

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

update_config_proto

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

update_config_proto(
    config_proto
)

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

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