tf.distribute.Strategy

Политика распределения состояния и вычислений по списку устройств.

tf.distribute.Strategy(
    extended
)

Обзор и примеры см. в руководстве. Глоссарий понятий, упомянутых на этой странице, таких как «на реплику», реплика и сведение, см. в разделах tf.distribute.StrategyExtended и tf.distribute.

Кратко:

• Чтобы использовать это с Keras compile/fit, ознакомьтесь с инструкцией.
• В противном случае используйте tf.distribute.Strategy.scope, чтобы указать, что при построении и выполнении модели следует использовать стратегию. (В результате вы окажетесь в «контексте между репликами» этой стратегии, то есть стратегия будет управлять такими аспектами, как размещение переменных.)

• Если вы пишете собственный цикл обучения, потребуется вызвать еще несколько методов; см. руководство:

  • Сначала создайте обычный объект tf.data.Dataset.
  • Используйте tf.distribute.Strategy.experimental_distribute_dataset, чтобы преобразовать tf.data.Dataset в объект, создающий значения «на реплику». Если вы хотите вручную указать, как набор данных должен распределяться между репликами, вместо этого используйте tf.distribute.Strategy.distribute_datasets_from_function.
  • Используйте tf.distribute.Strategy.run, чтобы выполнить функцию один раз для каждой реплики. Функция принимает значения, которые могут быть значениями «на реплику» (например, из объекта tf.distribute.DistributedDataset), и возвращает значения «на реплику». Эта функция выполняется в «контексте реплики», то есть каждая операция выполняется отдельно на каждой реплике.
  • Наконец, используйте метод (например, tf.distribute.Strategy.reduce), чтобы преобразовать полученные значения «на реплику» в обычные Tensor.

Собственный цикл обучения может быть очень простым:

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)

Этот код принимает обычные dataset и replica_fn и выполняет их распределенно с помощью указанной tf.distribute.Strategy, названной выше my_strategy. Все переменные, созданные в replica_fn, создаются с использованием политики my_strategy, а библиотечные функции, вызываемые replica_fn, могут использовать API get_replica_context() для реализации специфичного для распределенных вычислений поведения.

Вы можете использовать API reduce для агрегирования результатов по репликам и возвращать их как результат одной итерации по tf.distribute.DistributedDataset. Также можно использовать tf.keras.metrics (например, для потерь, точности и т. д.), чтобы накапливать метрики на протяжении шагов заданной эпохи.

Более подробный пример см. в руководстве по собственным циклам обучения.

Примечание: В настоящее время tf.distribute.Strategy не поддерживает секционированные переменные TensorFlow (когда одна переменная разделена между несколькими устройствами).

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