tf.distribute.Strategy

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

tf.distribute.Strategy(
    extended
)

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

Вкратце:

• Чтобы использовать его с Keras compile/fit, пожалуйста, ознакомьтесь.
• Вы можете передать потомка tf.distribute.Strategy в tf.estimator.RunConfig, чтобы указать, как tf.estimator.Estimator должен распределить свои вычисления. См. руководство.
• В противном случае, используйте 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 возвращенного этим 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 в новый размер пакета, равный глобальному размеру пакета, делённому на количество реплик в синхронизации. Мы перебираем его с помощью итератора 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>)

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

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-мя 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 перед выводом.

Примечание: Результат копируется на «текущее» устройство — обычно это ЦП узла, на котором выполняется программа. Для 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, поскольку они позволяют настраивать место назначения результата. Они также вызываются в контексте между репликами.

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

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

Например, если у вас есть глобальный размер пакета 8 и 2 реплики, значения для примеров [0, 1, 2, 3] будут на реплике 0, а [4, 5, 6, 7] — на реплике 1. С axis=None, reduce агрегирует только по репликам, возвращая [0+4, 1+5, 2+6, 3+7]. Это полезно, когда каждая реплика вычисляет скаляр или какое-либо другое значение без «размерности пакета» (например, градиент или потерю).

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

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

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

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

run

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

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

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

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

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

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

Что должно быть в контексте, а что — вне?

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

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