tf.distribute.Strategy

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

tf.distribute.Strategy(
    extended
)

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

Кратко:

• Для использования с 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 возвращаемого этим 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>)
       

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

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

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

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

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

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

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

reduce

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

reduce(
    reduce_op, value, axis
)

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

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
def step_fn():
  i = tf.distribute.get_replica_context().replica_id_in_sync_group
  return tf.identity(i)

per_replica_result = strategy.run(step_fn)
total = strategy.reduce("SUM", per_replica_result, axis=None)
total
<tf.Tensor: shape=(), dtype=int32, numpy=1>

Чтобы увидеть, как это будет выглядеть с несколькими репликами, рассмотрите тот же пример с MirroredStrategy с 2 графическими процессорами:

strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
def step_fn():
  i = tf.distribute.get_replica_context().replica_id_in_sync_group
  return tf.identity(i)

per_replica_result = strategy.run(step_fn)
# Check devices on which per replica result is:
strategy.experimental_local_results(per_replica_result)[0].device
# /job:localhost/replica:0/task:0/device:GPU:0
strategy.experimental_local_results(per_replica_result)[1].device
# /job:localhost/replica:0/task:0/device:GPU:1

total = strategy.reduce("SUM", per_replica_result, axis=None)
# Check device on which reduced result is:
total.device
# /job:localhost/replica:0/task:0/device:CPU:0

Этот API обычно используется для агрегирования результатов, возвращаемых разными репликами, для отчётности и т. п. Например, потерю, вычисленную на разных репликах, можно усреднить с помощью этого API перед печатью.

Примечание: Результат копируется на «текущее» устройство — обычно это процессор рабочего узла, на котором выполняется программа. Для TPUStrategy это первый хост TPU. Для многоклиентского MultiWorkerMirroredStrategy это процессор каждого рабочего узла.

Существует ряд различных API tf.distribute для сведения значений по репликам:

• tf.distribute.ReplicaContext.all_reduce: Это отличается от Strategy.reduce тем, что предназначено для контекста реплики и не копирует результаты на устройство хоста. all_reduce обычно используется для сокращений внутри шага обучения, таких как градиенты.
• tf.distribute.StrategyExtended.reduce_to и tf.distribute.StrategyExtended.batch_reduce_to: Эти API являются более продвинутыми версиями Strategy.reduce, поскольку они позволяют настроить место назначения результата. Они также вызываются в контексте межреплицированных вызовов.

Каким должен быть ось?

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

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

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

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

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

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

run

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

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

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

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

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

Все аргументы в args или kwargs могут быть вложенной структурой тензоров, например, списком тензоров, в этом случае args и kwargs будут переданы вызываемой fn на каждой реплике. Или args или kwargs могут быть tf.distribute.DistributedValues, содержащие тензоры или составные тензоры, т. е. tf.compat.v1.TensorInfo.CompositeTensor, в этом случае каждый вызов fn получит компонент tf.distribute.DistributedValues, соответствующий его реплике. Обратите внимание, что произвольные значения Python, которые не являются перечисленными выше типов, не поддерживаются.

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

1. Ввод постоянного тензора.

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   tensor_input = tf.constant(3.0)
   @tf.function
   def replica_fn(input):
     return input*2.0
   result = strategy.run(replica_fn, args=(tensor_input,))
   result
       PerReplica:{
         0: <tf.Tensor: shape=(), dtype=float32, numpy=6.0>,
         1: <tf.Tensor: shape=(), dtype=float32, numpy=6.0>
       }
       

2. Ввод DistributedValues.

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   @tf.function
   def run():
     def value_fn(value_context):
       return value_context.num_replicas_in_sync
     distributed_values = (
       strategy.experimental_distribute_values_from_function(
         value_fn))
     def replica_fn2(input):
       return input*2
     return strategy.run(replica_fn2, args=(distributed_values,))
   result = run()
   result
       <tf.Tensor: shape=(), dtype=int32, numpy=4>
       

3. Использование tf.distribute.ReplicaContext для allreduce значений.

   strategy = tf.distribute.MirroredStrategy(["gpu:0", "gpu:1"])
   @tf.function
   def run():
      def value_fn(value_context):
        return tf.constant(value_context.replica_id_in_sync_group)
      distributed_values = (
          strategy.experimental_distribute_values_from_function(
              value_fn))
      def replica_fn(input):
        return tf.distribute.get_replica_context().all_reduce(
            "sum", input)
      return strategy.run(replica_fn, args=(distributed_values,))
   result = run()
   result
       PerReplica:{
         0: <tf.Tensor: shape=(), dtype=int32, numpy=1>,
         1: <tf.Tensor: shape=(), dtype=int32, numpy=1>
       }
       

scope

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

scope()

Менеджер контекста для установки стратегии по умолчанию и распределения переменных.

Этот метод возвращает менеджер контекста и используется следующим образом:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
# Variable created inside scope:
with strategy.scope():
  mirrored_variable = tf.Variable(1.)
mirrored_variable
MirroredVariable:{
  0: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>,
  1: <tf.Variable 'Variable/replica_1:0' shape=() dtype=float32, numpy=1.0>
}
# Variable created outside scope:
regular_variable = tf.Variable(1.)
regular_variable
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>

Что происходит при входе в область Strategy.scope?

• strategy устанавливается в глобальный контекст как текущая стратегия. Внутри этой области tf.distribute.get_strategy() теперь вернет эту стратегию. За пределами этой области возвращается стратегия по умолчанию (бездействующая).
• Вход в область также означает вход в "меж-репличный контекст". Смотрите tf.distribute.StrategyExtended для объяснения меж-репличного и репличного контекстов.
• Создание переменных внутри scope перехватывается стратегией. Каждая стратегия определяет, как она должна повлиять на создание переменной. Синхронные стратегии, такие как MirroredStrategy, TPUStrategy и MultiWorkerMiroredStrategy, создают переменные, дублируемые на каждой реплике, а ParameterServerStrategy создаёт переменные на серверах параметров. Это выполняется с помощью настраиваемого tf.variable_creator_scope.
• В некоторых стратегиях может также быть вошла область по умолчанию для устройств: в MultiWorkerMiroredStrategy на каждом рабочем узле вводится область по умолчанию "/CPU:0".

Примечание: Вход в область не автоматически распределяет вычисление, за исключением случаев использования высокоуровневых фреймворков обучения, таких как keras model.fit. Если вы не используете model.fit, вам нужно использовать API strategy.run, чтобы явно распределить вычисление. Смотрите пример в учебнике по пользовательскому циклу обучения.

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

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

• Любой код, создающий переменные, которые должны быть распределёнными переменными, должен вызываться в области strategy.scope. Это можно сделать, либо непосредственно вызвав функцию создания переменных в контексте области, либо полагаясь на другой API, такой как strategy.run или keras.Model.fit, чтобы автоматически войти в него за вас. Любые переменные, созданные вне области, не будут распределены и могут привести к снижению производительности. Некоторые общие объекты, создающие переменные в TF, — это модели, оптимизаторы, метрики. Такие объекты всегда должны инициализироваться в области, и любые функции, которые могут лениво создавать переменные (например, Model.call(), отслеживания tf.function и т. д.), должны аналогичным образом вызываться в области. Другим источником создания переменных может быть восстановление контрольной точки — когда переменные создаются лениво. Обратите внимание, что любая переменная, созданная внутри стратегии, фиксирует информацию о стратегии. Таким образом, чтение и запись в эти переменные за пределами области strategy.scope также могут работать беспрепятственно, без необходимости ввода области пользователем.
• Некоторые API стратегии (например, strategy.run и strategy.reduce), которые должны быть в области стратегии, автоматически входят в область, что означает, что при использовании этих API вам не нужно явно входить в область.
• Когда tf.keras.Model создаётся внутри strategy.scope, объект модели фиксирует информацию об области. При вызове методов высокоуровневого обучения, таких как model.compile, model.fit и т. д., захваченная область будет автоматически входить, и связанная стратегия будет использоваться для распределения обучения и т. д. Подробный пример см. в учебнике по распределённому Keras. ПРЕДУПРЕЖДЕНИЕ: простой вызов model(..) не приводит к автоматическому входу в захваченную область — только высокоуровневые API для обучения поддерживают такое поведение: model.compile, model.fit, model.evaluate, model.predict и model.save могут вызываться как внутри, так и вне области.
• Следующие элементы могут быть как внутри, так и вне области:
  • Создание наборов данных входных данных
  • Определение tf.function представляющих шаг обучения
  • API сохранения, такие как tf.saved_model.save. Загрузка создаёт переменные, поэтому это должно быть в области, если вы хотите обучить модель распределённым способом.
  • Сохранение контрольных точек. Как указано выше — checkpoint.restore иногда может потребоваться быть внутри области, если оно создаёт переменные.