tf.compat.v1.distribute.MirroredStrategy

Синхронное обучение нескольких реплик на одном компьютере.

Наследует: Strategy

tf.compat.v1.distribute.MirroredStrategy(
    devices=None, cross_device_ops=None
)

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

Эта стратегия обычно используется для обучения на одном компьютере с несколькими GPU. Для TPU используйте tf.distribute.TPUStrategy. Сведения об использовании MirroredStrategy с несколькими рабочими узлами см. в разделе tf.distribute.experimental.MultiWorkerMirroredStrategy.

Например, переменная, созданная в MirroredStrategy, является MirroredVariable. Если в аргументе конструктора стратегии не указаны устройства, будут использоваться все доступные GPU. Если GPU не найдены, будут использоваться доступные CPU. Обратите внимание, что TensorFlow рассматривает все CPU компьютера как одно устройство и использует для распараллеливания внутренние потоки.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
with strategy.scope():
  x = tf.Variable(1.)
x
MirroredVariable:{
  0: <tf.Variable ... shape=() dtype=float32, numpy=1.0>,
  1: <tf.Variable ... shape=() dtype=float32, numpy=1.0>
}

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

Переменные, созданные внутри MirroredStrategy, обёрнутого в tf.function, по-прежнему являются MirroredVariables.

x = []
@tf.function  # Wrap the function with tf.function.
def create_variable():
  if not x:
    x.append(tf.Variable(1.))
  return x[0]
strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
with strategy.scope():
  _ = create_variable()
  print(x[0])
MirroredVariable:{
  0: <tf.Variable ... shape=() dtype=float32, numpy=1.0>,
  1: <tf.Variable ... shape=() dtype=float32, numpy=1.0>
}

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

Например:

my_strategy = tf.distribute.MirroredStrategy()
with my_strategy.scope():
  @tf.function
  def distribute_train_epoch(dataset):
    def replica_fn(input):
      # process input and return result
      return result

    total_result = 0
    for x in dataset:
      per_replica_result = my_strategy.run(replica_fn, args=(x,))
      total_result += my_strategy.reduce(tf.distribute.ReduceOp.SUM,
                                         per_replica_result, axis=None)
    return total_result

  dist_dataset = my_strategy.experimental_distribute_dataset(dataset)
  for _ in range(EPOCHS):
    train_result = distribute_train_epoch(dist_dataset)

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 execution выполняет операции, заданные в 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)

Возвращаемый функцией input_fn объект tf.data.Dataset должен иметь размер пакета на реплику; его можно вычислить с помощью 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: эти API являются более продвинутыми вариантами Strategy.reduce, поскольку позволяют настраивать устройство назначения результата. Они также вызываются в контексте между репликами.

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

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

run

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

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

Вызывает fn на каждой реплике с заданными аргументами.

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

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

Все аргументы в args или kwargs могут представлять собой вложенную структуру тензоров, например список тензоров. В таком случае args и kwargs будут переданы в fn, вызываемый на каждой реплике. Кроме того, args или kwargs могут быть объектами tf.distribute.DistributedValues, содержащими тензоры или составные тензоры, то есть tf.compat.v1.TensorInfo.CompositeTensor. В этом случае каждый вызов fn получит компонент tf.distribute.DistributedValues, соответствующий его реплике. Обратите внимание, что произвольные значения Python, не относящиеся к перечисленным выше типам, не поддерживаются.

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

1. Входной тензор-константа.

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   tensor_input = tf.constant(3.0)
   @tf.function
   def replica_fn(input):
     return input*2.0
   result = strategy.run(replica_fn, args=(tensor_input,))
   result
       PerReplica:{
         0: <tf.Tensor: shape=(), dtype=float32, numpy=6.0>,
         1: <tf.Tensor: shape=(), dtype=float32, numpy=6.0>
       }
       

2. Входные данные DistributedValues.

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   @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_fn2(input):
       return input*2
     return strategy.run(replica_fn2, args=(distributed_values,))
   result = run()
   result
       <tf.Tensor: shape=(), dtype=int32, numpy=4>
       

3. Использование tf.distribute.ReplicaContext для выполнения all-reduce значений.

   strategy = tf.distribute.MirroredStrategy(["gpu:0", "gpu:1"])
   @tf.function
   def run():
      def value_fn(value_context):
        return tf.constant(value_context.replica_id_in_sync_group)
      distributed_values = (
          strategy.experimental_distribute_values_from_function(
              value_fn))
      def replica_fn(input):
        return tf.distribute.get_replica_context().all_reduce(
            "sum", input)
      return strategy.run(replica_fn, args=(distributed_values,))
   result = run()
   result
       PerReplica:{
         0: <tf.Tensor: shape=(), dtype=int32, numpy=1>,
         1: <tf.Tensor: shape=(), dtype=int32, numpy=1>
       }
       

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. Пример см. в руководстве по пользовательскому циклу обучения.

Что должно находиться внутри области, а что — вне её?

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

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

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