tf.distribute.MirroredStrategy

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

Наследует: Strategy

tf.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_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(["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. Распределение значений в массиве на основе replica_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, каждый рабочий процесс является отдельным клиентом, и эта функция возвращает только значения, вычисленные на этом рабочем процессе.

gather

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

gather(
    value, axis
)

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

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

Этот 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)>

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 и двумя 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, а реплик две, значения для примеров [0, 1, 2, 3] будут находиться на реплике 0, а [4, 5, 6, 7] — на реплике 1. При axis=None операция reduce будет выполнять агрегацию только между репликами и возвращать [0+4, 1+5, 2+6, 3+7]. Это полезно, если каждая реплика вычисляет скаляр или какое-либо другое значение без измерения «пакета» (например, градиент или потери).

strategy.reduce("sum", per_replica_result, axis=None)

Иногда требуется агрегировать данные как по всему глобальному пакету, так и по всем репликам. Такое поведение можно получить, указав измерение пакета как axis, обычно axis=0. В этом случае будет возвращен скаляр 0+1+2+3+4+5+6+7.

strategy.reduce("sum", per_replica_result, axis=0)

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

run

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

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

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

Этот метод является основным способом распределения вычислений с помощью объекта 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 для выполнения allreduce значений.

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

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

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

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