tf.distribute.experimental.TPUStrategy
Синхронное обучение на TPUs и TPU Pod.
Наследуется от: Strategy
tf.distribute.experimental.TPUStrategy(
tpu_cluster_resolver=None, device_assignment=None
)
Используется в ноутбуках
| Используется в руководствах |
|---|
Для создания объекта 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 | Возвращает решатель кластера, связанный с этой стратегией.
|
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 для примера.
Примечание: Если вы используете 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, содержащий значение для каждой реплики. |
Пример использования:
-
Возвращение константного значения для каждой реплики:
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>) -
Распределение значений в массиве на основе 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) -
Указание значений с помощью 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) -
Размещение значений на устройствах и их распределение:
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.DistributedValuesvalue, его компоненты тензоров должны иметь ненулевой ранг. В противном случае рассмотрите использование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» вводится на каждом рабочем узле.
Примечание: Вхождение в контекст не автоматически распределяет вычисление, за исключением случаев использования высокоуровневых фреймворков обучения, таких как kerasmodel.fit. Если вы не используетеmodel.fit, вам необходимо использоватьstrategy.runAPI для явного распределения этого вычисления. См. пример в учебном пособии по настраиваемым циклам обучения 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