tf.distribute.MirroredStrategy

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

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

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

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

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

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

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 будет вызвана на ЦП устройства каждого из рабочих узлов, и каждый из них сгенерирует набор данных, где каждая реплика на этом рабочем узле будет извлекать одну партию входных данных (т. е. если на рабочем узле две реплики, с каждым шагом будут извлечены две партии из 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.

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

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 по-питоновски. 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. Распределение значений в массиве на основе идентификатора реплики:

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

Этот API можно вызывать только в контексте между репликами. Для аналога в контексте реплики см. tf.distribute.ReplicaContext.all_gather.

Примечание: Для всех стратегий, кроме tf.distribute.TPUStrategy, вход value на разных репликах должен иметь одинаковый ранг, а их формы должны быть одинаковыми во всех измерениях, кроме измерения axis. Другими словами, их формы не могут отличаться в измерении d, где d не равно аргументу axis. Например, для tf.distribute.DistributedValues с тензорами компонентов формы (1, 2, 3) и (1, 3, 3) на двух репликах, вы можете вызвать gather(..., axis=1, ...), но не gather(..., axis=0, ...) или gather(..., axis=2, ...). Однако для tf.distribute.TPUStrategy.gather все тензоры должны иметь точно одинаковый ранг и одинаковую форму.

Примечание: Для tf.distribute.DistributedValues value тензоры компонентов должны иметь ненулевой ранг. В противном случае рассмотрите использование tf.expand_dims перед их сборкой.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
# A DistributedValues with component tensor of shape (2, 1) on each replica
distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(tf.constant([[1], [2]])))
@tf.function
def run():
  return strategy.gather(distributed_values, axis=0)
run()
<tf.Tensor: shape=(4, 1), dtype=int32, numpy=
array([[1],
       [2],
       [1],
       [2]], dtype=int32)>

Рассмотрим следующий пример для других комбинаций:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1", "GPU:2", "GPU:3"])
single_tensor = tf.reshape(tf.range(6), shape=(1,2,3))
distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(single_tensor))
@tf.function
def run(axis):
  return strategy.gather(distributed_values, axis=axis)
axis=0
run(axis)
<tf.Tensor: shape=(4, 2, 3), dtype=int32, numpy=
array([[[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]],
       [[0, 1, 2],
        [3, 4, 5]]], dtype=int32)>
axis=1
run(axis)
<tf.Tensor: shape=(1, 8, 3), dtype=int32, numpy=
array([[[0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5],
        [0, 1, 2],
        [3, 4, 5]]], dtype=int32)>
axis=2
run(axis)
<tf.Tensor: shape=(1, 2, 12), dtype=int32, numpy=
array([[[0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2],
        [3, 4, 5, 3, 4, 5, 3, 4, 5, 3, 4, 5]]], dtype=int32)>

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

Что должно быть осью?

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

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