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 не найдено, она будет использовать доступные 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 возвращаемого этим 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>)

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

1. Указание значений с помощью 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)

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

gather

Просмотр исходного кода

gather(
    value, axis
)

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

Учитывая tf.distribute.DistributedValues или подобный объекту tf.Tensor value, эта API собирает и конкатенирует value по репликам вдоль axis-й размерности. Результат копируется на "текущее" устройство

• которое, как правило, является процессором рабочего узла, на котором выполняется программа. Для tf.distribute.TPUStrategy это первый хост TPU. Для многоклиентских 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 должны быть либо Python-значениями вложенной структуры тензоров, например, списком тензоров, в этом случае args и kwargs будут переданы в fn , вызываемый на каждой реплике. Или args или kwargs могут быть tf.distribute.DistributedValues, содержащими тензоры или составные тензоры, т. е. tf.compat.v1.TensorInfo.CompositeTensor, в этом случае каждый вызов fn получит компонент tf.distribute.DistributedValues, соответствующий его реплике.

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

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>
}

1. Входные значения 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>

1. Использование 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 для явного распределения вычислений. См. пример в руководстве по пользовательскому циклу обучения https://www.tensorflow.org/tutorials/distribute/custom_training.

Что должно быть в области действия, а что вне ее?

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

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