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 до нового размера пакета, равного глобальному размеру пакета, делённому на количество реплик в синхронизации. Мы перебираем его с помощью цикла 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>)

1. Распределение значений в массиве на основе идентификатора реплики:

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

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