tf.distribute.TPUStrategy

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

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

tf.distribute.TPUStrategy(
    tpu_cluster_resolver=None, experimental_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.TPUStrategy(resolver)

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

Для запуска программ TF2 на TPUs можно использовать 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.

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

Пример написания настраиваемой процедуры обучения на TPUs:

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.experimental_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)

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

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.experimental_distribute_datasets_from_function(
    dataset_fn)
iterator = iter(dist_dataset)
strategy.run(step_fn, args=(next(iterator),))

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.

Методы

experimental_assign_to_logical_device

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

experimental_assign_to_logical_device(
    tensor, logical_device_id
)

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

Примечание: Этот API в настоящее время поддерживается только в TPUStrategy. Это добавляет аннотацию к 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.

Следующий пример:

strategy = tf.distribute.MirroredStrategy()

# Create a dataset
dataset = dataset_ops.Dataset.TFRecordDataset([
  "/a/1.tfr", "/a/2.tfr", "/a/3.tfr", "/a/4.tfr"])

# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(dataset)

# Iterate over the `tf.distribute.DistributedDataset`
for x in dist_dataset:
  # process dataset elements
  strategy.run(replica_fn, args=(x,))

В приведённом фрагменте кода tf.distribute.DistributedDataset dist_dataset сгруппирован по GLOBAL_BATCH_SIZE, и мы перебираем его с помощью for x in dist_dataset. x tf.distribute.DistributedValues, содержащий данные для всех реплик, агрегируются в пакет из GLOBAL_BATCH_SIZE. tf.distribute.Strategy.run позаботится о предоставлении правильных данных для каждой реплики в x в соответствующую replica_fn , выполняемую на каждой реплике.

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

Если попытка разбить по файлам не удалась (т.е. набор данных не считывается из файлов), мы разделим набор данных равномерно в конце, добавив операцию .shard в конец конвейера обработки. Это заставит весь конвейер предобработки данных быть выполнен на каждом узле, и каждый узел выполнит избыточную работу. Мы выведем предупреждение, если этот путь будет выбран.

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

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

Вы также можете использовать свойство element_spec экземпляра tf.distribute.DistributedDataset, возвращённого этим API, для запроса tf.TypeSpec элементов, возвращаемых итератором. Это можно использовать для установки свойства input_signature tf.function.

strategy = tf.distribute.MirroredStrategy()

# Create a dataset
dataset = dataset_ops.Dataset.TFRecordDataset([
  "/a/1.tfr", "/a/2.tfr", "/a/3.tfr", "/a/4.tfr"])

# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(dataset)

@tf.function(input_signature=[dist_dataset.element_spec])
def train_step(inputs):
  # train model with inputs
  return

# Iterate over the `tf.distribute.DistributedDataset`
for x in dist_dataset:
  # process dataset elements
  strategy.run(train_step, args=(x,))

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

experimental_distribute_datasets_from_function

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

experimental_distribute_datasets_from_function(
    dataset_fn, options=None
)

Распределяет экземпляры tf.data.Dataset, созданные вызовами dataset_fn.

dataset_fn будет вызван один раз для каждого работника в стратегии. Каждый экземпляр на этом работнике будет извлекать одну партию входов из локального Dataset (т. е. если у работника есть два экземпляра, две партии будут извлечены из Dataset на каждом шаге).

Этот метод может использоваться для нескольких целей. Например, там, где 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.

global_batch_size = 8
def dataset_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)

strategy = tf.distribute.MirroredStrategy()
ds = strategy.experimental_distribute_datasets_from_function(dataset_fn)

def train(ds):
  @tf.function(input_signature=[ds.element_spec])
  def step_fn(inputs):
    # train the model with inputs
    return inputs

... for batch in ds: ... replica_results = strategy.run(replica_fn, args=(batch,))

train(ds)

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

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()
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>,)

1. Распределение значений в массиве на основе идентификатора реплики:

strategy = tf.distribute.MirroredStrategy()
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,)

1. Указание значений с использованием num_replicas_in_sync:

strategy = tf.distribute.MirroredStrategy()
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
(1,)

1. Размещение значений на устройствах и распределение:

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_make_numpy_dataset

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

experimental_make_numpy_dataset(
    numpy_input
)

Создаёт tf.data.Dataset из массива NumPy. (устарело)

Это позволяет избежать добавления numpy_input как большого константы в граф и копирует данные на машину или машины, которые будут обрабатывать ввод.

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

Пример:

strategy = tf.distribute.MirroredStrategy()
numpy_input = np.ones([10], dtype=np.float32)
dataset = strategy.experimental_make_numpy_dataset(numpy_input)
dataset
<TensorSliceDataset shapes: (), types: tf.float32>
dataset = dataset.batch(2)
dist_dataset = strategy.experimental_distribute_dataset(dataset)

experimental_replicate_to_logical_devices

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

experimental_replicate_to_logical_devices(
    tensor
)

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

Примечание: Этот API поддерживается только в TPUStrategy на данный момент. Это добавляет аннотацию к тензору 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),))

Аргументы: тензор: входной тензор для аннотации.

experimental_split_to_logical_devices

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

experimental_split_to_logical_devices(
    tensor, partition_dimensions
)

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

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

Например, для системы с 8 логическими устройствами, если tensor — это тензор изображения с формой (batch_size, width, height, channel) и partition_dimensions — [1, 2, 4, 1], тогда tensor будет разделён на 2 в ширину и 4 в высоту, и значения разделённого тензора будут переданы на 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)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)

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),))

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

reduce

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

reduce(
    reduce_op, value, axis
)

Сведение value по репликам.

Учитывая значение по реплике, возвращаемое run, например, потерю на пример, пакет будет разделен между всеми репликами. Эта функция позволяет вам агрегировать по репликам и, необязательно, также по элементам пакета. Например, если у вас есть глобальный размер пакета 8 и 2 реплики, значения для примеров [0, 1, 2, 3] будут на реплике 0, а [4, 5, 6, 7] — на реплике 1. По умолчанию reduce будет просто агрегировать по репликам, возвращая [0+4, 1+5, 2+6, 3+7]. Это полезно, когда каждая реплика вычисляет скаляр или какое-либо другое значение без измерения "пакет" (например, градиент). Чаще вы захотите агрегировать по глобальному пакету, что можно сделать, указав измерение пакета как axis, обычно axis=0. В этом случае она вернет скаляр 0+1+2+3+4+5+6+7.

Если есть последний частичный пакет, вам нужно будет указать ось, чтобы результирующая форма была согласована между репликами. Итак, если последний пакет имеет размер 6 и он разделен на [0, 1, 2, 3] и [4, 5], у вас возникнет несоответствие форм, если вы не укажете axis=0. Если вы укажете tf.distribute.ReduceOp.MEAN, используя axis=0 будет использоваться правитель один-шесть. Противопоставьте это вычислению 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.experimental_distribute_datasets_from_function, когда fn выполняется на определенной реплике, она будет выполнена с компонентом tf.distribute.DistributedValues, соответствующим этой реплике.

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

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

Example usage:

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()
# 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>
}
# 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 или model.fit, для входа в него за вас. Любая переменная, созданная за пределами области, не будет распределена и может иметь последствия для производительности. Обычные вещи, которые создают переменные в TF: модели, оптимизаторы, метрики. Они всегда должны создаваться внутри области. Другим источником создания переменных может быть восстановление контрольной точки - когда переменные создаются лениво. Обратите внимание, что любая переменная, созданная внутри стратегии, захватывает информацию о стратегии. Таким образом, чтение и запись в эти переменные вне strategy.scope также могут работать беспрепятственно, без необходимости входа пользователя в область.
• Некоторые API стратегий (такие как strategy.run и strategy.reduce), которые требуют нахождения в области стратегии, автоматически входят в область, что означает, что при использовании этих API вам не нужно входить в область самостоятельно.
• Когда tf.keras.Model создается внутри strategy.scope, мы захватываем эту информацию. Когда методы высокоуровневых фреймворков обучения, такие как 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 иногда может потребоваться быть внутри области, если это создает переменные.