Spec-Zone.ru › TensorFlow

tf.distribute.experimental.TPUStrategy

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

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

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

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

Используется в руководствах
  • Классификация изображений с помощью Model Garden
  • Сегментация инстансов с помощью Model Garden
  • Обнаружение объектов с помощью Model Garden
  • Семантическая сегментация с помощью Model Garden

Для создания объекта 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.experimental.TPUStrategy(resolver)

При использовании стратегий распределения переменные, созданные в области видимости стратегии, будут дублированы на всех репликах и могут быть синхронизированы с помощью алгоритмов all-reduce.

Для запуска программ TF2 на TPUs можно использовать .compile и .fit API в tf.keras вместе с TPUStrategy или написать собственную настраиваемую петлю обучения, вызвав strategy.run непосредственно. Обратите внимание, что TPUStrategy не поддерживает чистое жадное выполнение, поэтому убедитесь, что функция, переданная в strategy.run, является tf.function или strategy.run вызывается внутри tf.function, если включено поведение жадной оценки.

Args
tpu_cluster_resolver tf.distribute.cluster_resolver.TPUClusterResolver, предоставляющий информацию о кластере TPU.
device_assignment Необязательный tf.tpu.experimental.DeviceAssignment для указания размещения реплик в кластере TPU.
Атрибуты
cluster_resolver Возвращает решатель кластера, связанный с этой стратегией.

tf.distribute.experimental.TPUStrategy предоставляет связанный tf.distribute.cluster_resolver.ClusterResolver. Если пользователь предоставляет его в __init__, возвращается этот экземпляр; если нет, предоставляется по умолчанию tf.distribute.cluster_resolver.TPUClusterResolver.

extended tf.distribute.StrategyExtended с дополнительными методами.
num_replicas_in_sync Возвращает количество реплик, по которым агрегируются градиенты.

Методы

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 для примера.

Важно: Набор данных tf.data.Dataset, возвращаемый dataset_fn, должен иметь размер пакета на реплику, в отличие от experimental_distribute_dataset, который использует общий размер пакета. Это можно вычислить с помощью input_context.get_per_replica_batch_size.
Примечание: Если вы используете 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.

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

Args
dataset_fn Функция, принимающая экземпляр tf.distribute.InputContext и возвращающая экземпляр tf.data.Dataset.
options tf.distribute.InputOptions для управления параметрами распределения этого набора данных.
Возвращает
tf.distribute.DistributedDataset.

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

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

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

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

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

Аргументы
dataset tf.data.Dataset, который будет разделен между всеми репликами по указанным выше правилам.
options tf.distribute.InputOptions, используемые для управления параметрами распределения этого набора данных.
Возвращаемое значение
tf.distribute.DistributedDataset.

experimental_distribute_values_from_function

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

experimental_distribute_values_from_function(
    value_fn
)

Генерирует tf.distribute.DistributedValues из value_fn.

Эта функция предназначена для генерации tf.distribute.DistributedValues для передачи в run, reduce или другие методы, принимающие распределённые значения, когда не используются наборы данных.

Аргументы
value_fn Функция для выполнения, чтобы сгенерировать значения. Она вызывается для каждой реплики с tf.distribute.ValueContext в качестве единственного аргумента. Она должна возвращать тензор или тип, который можно преобразовать в тензор.
Возвращаемое значение
tf.distribute.DistributedValues, содержащий значение для каждой реплики.

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

  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. Распределение значений в массиве на основе 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, каждый рабочий процесс будет своим клиентом, и эта функция вернёт только значения, вычисленные на этом рабочем процессе.
Аргументы
value Значение, возвращённое experimental_run(), run(), or a variable created inscope`.
Возвращаемое значение
Кортеж значений, содержащихся в value, где i-й элемент соответствует i-й реплике. Если value представляет единственное значение, это возвращает (value,).

gather

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

gather(
    value, axis
)

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

Учитывая tf.distribute.DistributedValues или tf.Tensor-подобный объект value, эта API собирает и конкатенирует value по репликам вдоль axis-й размерности. Результат копируется на «текущее» устройство, которое обычно является процессором рабочего процесса, на котором выполняется программа. Для tf.distribute.TPUStrategy это первый хост TPU. Для многоклиентского tf.distribute.MultiWorkerMirroredStrategy это процессор каждого рабочего процесса.

Эта 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)>
Аргументы
value экземпляр tf.distribute.DistributedValues, например, возвращённый Strategy.run, для объединения в один тензор. Он также может быть обычным тензором при использовании с tf.distribute.OneDeviceStrategy или по умолчанию. Тензоры, составляющие DistributedValues, могут быть только плоскими тензорами с ненулевым рангом, а не tf.IndexedSlices.
axis 0-мерный тензор int32. Размерность для объединения. Должно быть в диапазоне [0, ранг(значение)).
Возвращаемое значение
A Tensor, представляющий собой конкатенацию value по всем репликам вдоль axis измерения.

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

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

Учитывая значение, возвращаемое каждой репликой run, скажем, потерю на пример, пакет будет разделен между всеми репликами. Эта функция позволяет вам агрегировать по репликам и, по желанию, также по элементам пакета, указав параметр axis соответствующим образом.

Например, если у вас есть общий размер пакета 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.

Аргументы
reduce_op значение tf.distribute.ReduceOp, указывающее, как следует объединять значения. Разрешает использовать строковое представление перечисления, такое как «SUM», «MEAN».
value экземпляр tf.distribute.DistributedValues, например, возвращаемый Strategy.run, для объединения в один тензор. Он также может быть обычным тензором, если используется с OneDeviceStrategy или стратегией по умолчанию.
axis указывает размерность для редукции вдоль тензора каждой реплики. Обычно следует устанавливать в размерность пакета или None, чтобы уменьшить только по репликам (например, если у тензора нет размера пакета).
Возвращаемое значение
A Tensor.

run

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

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

См. базовый класс.

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

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

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

  • Любые операции, создающие переменные, которые должны быть распределёнными переменными, должны вызываться в области видимости strategy.scope. Это можно сделать, либо напрямую вызвав функцию создания переменной внутри контекста области видимости, либо, положившись на другой API, такой как strategy.run или keras.Model.fit, чтобы он автоматически вошел в него за вас. Любая переменная, созданная вне области видимости, не будет распределена и может иметь последствия для производительности. Некоторые общие объекты, создающие переменные в TF, — это модели, оптимизаторы, метрики. Такие объекты всегда должны инициализироваться в области видимости, а любые функции, которые могут создавать переменные лениво (например, Model.call(), отслеживание tf.function и т. п.), аналогично должны вызываться в пределах области видимости. Ещё один источник создания переменных — это восстановление контрольной точки — когда переменные создаются лениво. Обратите внимание, что любая переменная, созданная внутри стратегии, сохраняет информацию о стратегии. Поэтому чтение и запись этих переменных вне области видимости strategy.scope также могут работать без проблем, без необходимости вхождения пользователя в область видимости.
  • Некоторые API стратегий (например, strategy.run и strategy.reduce), которые требуют нахождения в области видимости стратегии, автоматически входят в эту область видимости, что означает, что при использовании этих API вам не нужно явно входить в неё.
  • Когда tf.keras.Model создаётся внутри strategy.scope, объект Model сохраняет информацию о области видимости. Когда вызываются методы высокоуровневых фреймворков обучения, такие как model.compile, model.fit и т. д., захваченная область видимости автоматически будет входить в неё, и связанная стратегия будет использоваться для распределения обучения и т. д. Подробный пример см. в учебном пособии по распределённому Keras. ВНИМАНИЕ: просто вызов model(..) не автоматически входит в захваченную область видимости — только высокоуровневые API обучения поддерживают такое поведение: model.compile, model.fit, model.evaluate, model.predict и model.save могут быть вызваны внутри или вне области видимости.
  • Следующее может быть как внутри, так и вне области видимости:
    • Создание наборов данных входных данных
    • Определение tf.function которые представляют ваш шаг обучения
    • API сохранения, такие как tf.saved_model.save. Загрузка создаёт переменные, поэтому это должно происходить внутри области видимости, если вы хотите обучить модель распределённым способом.
    • Сохранение контрольных точек. Как упоминалось выше — checkpoint.restore иногда может потребоваться внутри области видимости, если он создаёт переменные.
Возвращаемое значение
Менеджер контекста.

© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/distribute/experimental/TPUStrategy

Spec-Zone.ru

Настройки Оффлайн Что нового Помощь О нас
Spec-Zone .ru
спецификации, руководства, описания, API