tf.compat.v1.distribute.MirroredStrategy

Синхронное обучение на нескольких репликах на одной машине.

Наследуется от: Strategy

tf.compat.v1.distribute.MirroredStrategy(
    devices=None, cross_device_ops=None
)

Эта стратегия обычно используется для обучения на одной машине с несколькими графическими процессорами. Для TPUs используйте tf.distribute.TPUStrategy. Чтобы использовать MirroredStrategy с несколькими рабочими узлами, обратитесь к tf.distribute.experimental.MultiWorkerMirroredStrategy.

Например, переменная, созданная в рамках MirroredStrategy , является MirroredVariable. Если устройства не указаны в аргументе конструктора стратегии, она будет использовать все доступные графические процессоры. Если графических процессоров не найдено, будет использован доступный процессор ЦП. Обратите внимание, что TensorFlow обрабатывает все процессоры ЦП на машине как одно устройство и использует потоки для параллелизма внутри.

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
with strategy.scope():
  x = tf.Variable(1.)
x
MirroredVariable:{
  0: <tf.Variable ... shape=() dtype=float32, numpy=1.0>,
  1: <tf.Variable ... shape=() dtype=float32, numpy=1.0>
}

При использовании стратегий распределения все создание переменных должно выполняться в рамках области действия стратегии. Это позволит реплицировать переменные по всем репликам и поддерживать их синхронизацию с помощью алгоритма all-reduce.

Переменные, созданные внутри MirroredStrategy , которая обернута функцией tf.function, всё ещё MirroredVariables.

x = []
@tf.function  # Wrap the function with tf.function.
def create_variable():
  if not x:
    x.append(tf.Variable(1.))
  return x[0]
strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
with strategy.scope():
  _ = create_variable()
  print(x[0])
MirroredVariable:{
  0: <tf.Variable ... shape=() dtype=float32, numpy=1.0>,
  1: <tf.Variable ... shape=() dtype=float32, numpy=1.0>
}

experimental_distribute_dataset можно использовать для распределения набора данных по репликам при написании собственного цикла обучения. Если вы используете .fit и .compile методы, доступные в tf.keras, то tf.keras будет обрабатывать распределение за вас.

Например:

my_strategy = tf.distribute.MirroredStrategy()
with my_strategy.scope():
  @tf.function
  def distribute_train_epoch(dataset):
    def replica_fn(input):
      # process input and return result
      return result

    total_result = 0
    for x in dataset:
      per_replica_result = my_strategy.run(replica_fn, args=(x,))
      total_result += my_strategy.reduce(tf.distribute.ReduceOp.SUM,
                                         per_replica_result, axis=None)
    return total_result

  dist_dataset = my_strategy.experimental_distribute_dataset(dataset)
  for _ in range(EPOCHS):
    train_result = distribute_train_epoch(dist_dataset)

os.environ['TF_CONFIG'] = json.dumps({
  'cluster': {
      'worker': ["localhost:12345", "localhost:23456"],
      'ps': ["localhost:34567"]
  },
  'task': {'type': 'worker', 'index': 0}
})

# This implicitly uses TF_CONFIG for the cluster and current task info.
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()

...

if strategy.cluster_resolver.task_type == 'worker':
  # Perform something that's only applicable on workers. Since we set this
  # as a worker above, this block will run on this particular instance.
elif strategy.cluster_resolver.task_type == 'ps':
  # Perform something that's only applicable on parameter servers. Since we
  # set this as a worker above, this block will not run on this particular
  # instance.

Методы

distribute_datasets_from_function

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

distribute_datasets_from_function(
    dataset_fn, options=None
)

Распределяет экземпляры tf.data.Dataset, созданные вызовами dataset_fn.

Аргумент dataset_fn , который передается пользователем, — это функция-вход, имеющая аргумент tf.distribute.InputContext и возвращающая экземпляр tf.data.Dataset. Ожидается, что возвращаемый из dataset_fn набор данных уже сгруппирован по размерам пакета на реплику (то есть, глобальный размер пакета, деленный на количество реплик, синхронизированных в паре) и разделен. tf.distribute.Strategy.distribute_datasets_from_function не группирует и не разделяет экземпляр tf.data.Dataset, возвращаемый функцией-входом. dataset_fn будет вызываться на устройстве ЦП каждого из рабочих узлов, и каждый создаёт набор данных, где каждая реплика на этом рабочем узле будет декодировать один пакет входных данных (то есть, если у рабочего узла две реплики, две порции будут декодированы из Dataset на каждом шаге).

Этот метод можно использовать для нескольких целей. Во-первых, он позволяет задавать собственную логику группирования и разделения. (В отличие от tf.distribute.experimental_distribute_dataset, которая выполняет группировку и разделение за вас.) Например, там, где experimental_distribute_dataset не может разделить входные файлы, этот метод может быть использован для ручного разделения набора данных (избегая медленной обратной совместимости в experimental_distribute_dataset). В случаях, когда набор данных бесконечный, это разделение можно выполнить, создав реплики набора данных, которые отличаются только случайным началом.

Функция dataset_fn должна принимать экземпляр tf.distribute.InputContext, где можно получить доступ к информации о группировании и репликации входных данных.

Вы можете использовать свойство element_spec возвращённого этим API tf.distribute.DistributedDataset, чтобы получить доступ к tf.TypeSpec элементов, возвращаемых итератором. Это может быть использовано для задания свойства input_signature объекта tf.function. Последовайте tf.distribute.DistributedDataset.element_spec, чтобы посмотреть пример.

Примечание: Если вы используете TPUStrategy, порядок обработки данных рабочими узлами при использовании tf.distribute.Strategy.experimental_distribute_dataset или tf.distribute.Strategy.distribute_datasets_from_function не гарантируется. Это обычно требуется, если вы используете tf.distribute для масштабирования предсказаний. Однако вы можете вставить индекс для каждого элемента в пакете и упорядочить результаты соответствующим образом. Обратитесь к этому фрагменту для примера упорядочения результатов.

Примечание: Состоятельные преобразования набора данных в настоящее время не поддерживаются с tf.distribute.experimental_distribute_dataset или tf.distribute.distribute_datasets_from_function. Любые состоятельные операции, которые может иметь набор данных, в настоящее время игнорируются. Например, если ваш набор данных имеет преобразование map_fn , которое использует tf.random.uniform для поворота изображения, то у вас есть граф набора данных, который зависит от состояния (то есть от начального случайного числа) на локальной машине, где выполняется процесс Python.

Для получения дополнительной информации и свойств этого метода обратитесь к учебнику по распределенным входам. Если вы заинтересованы в обработке последних неполных пакетов, прочтите эту секцию.

experimental_distribute_dataset

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

experimental_distribute_dataset(
    dataset, options=None
)

Создаёт tf.distribute.DistributedDataset из tf.data.Dataset.

Возвращаемый tf.distribute.DistributedDataset можно перебирать аналогично обычным наборам данных. ПРИМЕЧАНИЕ: пользователь не может добавить больше преобразований к tf.distribute.DistributedDataset. Вы можете только создать итератор или изучить tf.TypeSpec генерируемых данных. См. документацию API tf.distribute.DistributedDataset для получения дополнительной информации.

Следующий пример:

global_batch_size = 2
# Passing the devices is optional.
strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
# Create a dataset
dataset = tf.data.Dataset.range(4).batch(global_batch_size)
# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(dataset)
@tf.function
def replica_fn(input):
  return input*2
result = []
# Iterate over the `tf.distribute.DistributedDataset`
for x in dist_dataset:
  # process dataset elements
  result.append(strategy.run(replica_fn, args=(x,)))
print(result)
[PerReplica:{
  0: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([0])>,
  1: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([2])>
}, PerReplica:{
  0: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([4])>,
  1: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([6])>
}]

Три ключевых действия, происходящих в этом методе: группировка, разделение и предварительная выборка.

В приведенном выше фрагменте кода dataset сгруппирован по global_batch_size, и вызов experimental_distribute_dataset для него перегруппирует dataset в новую размерность пакета, равную глобальной размерности пакета, деленной на количество реплик в синхронизации. Мы проходим по нему с помощью питоновского цикла for. x — это tf.distribute.DistributedValues, содержащий данные для всех реплик, и каждая реплика получает данные новой размерности пакета. tf.distribute.Strategy.run позаботится о предоставлении правильных данных на реплику в x для правильного replica_fn выполнения на каждой реплике.

Разбиение содержит автоматическое разбиение по множеству рабочих процессов и внутри каждого рабочего процесса. Во-первых, в распределённом обучении по нескольким рабочим процессам (т.е. когда вы используете tf.distribute.experimental.MultiWorkerMirroredStrategy или tf.distribute.TPUStrategy), автоматическое разбиение набора данных по нескольким рабочим процессам означает, что каждому рабочему процессу назначается подмножество всего набора данных (если установлен правильный tf.data.experimental.AutoShardPolicy). Это делается для того, чтобы на каждом шаге глобальный размер пакета непересекающихся элементов набора данных обрабатывался каждым рабочим процессом. Автоматическое разбиение имеет несколько вариантов, которые можно указать с помощью tf.data.experimental.DistributeOptions. Затем, разбиение внутри каждого рабочего процесса означает, что метод разделит данные между всеми устройствами рабочего процесса (если их несколько). Это произойдёт независимо от автоматического разбиения по нескольким рабочим процессам.

Примечание: для автоматического разбиения по нескольким рабочим процессам, режим по умолчанию — tf.data.experimental.AutoShardPolicy.AUTO. Этот режим попытается разбить входной набор данных по файлам, если набор данных создаётся из наборов данных для чтения (например, tf.data.TFRecordDataset, tf.data.TextLineDataset и т.д.) или в противном случае разбить набор данных по данным, где каждый из рабочих процессов будет читать весь набор данных и обрабатывать только присвоенную ему часть. Однако, если у вас меньше одного входного файла на рабочий процесс, рекомендуется отключить автоматическое разбиение наборов данных по рабочим процессам, установив tf.data.experimental.DistributeOptions.auto_shard_policy в tf.data.experimental.AutoShardPolicy.OFF.

По умолчанию этот метод добавляет преобразование предварительной выборки в конце предоставленного пользователем экземпляра tf.data.Dataset. Аргумент преобразования предварительной выборки, который buffer_size равен числу реплик в синхронизации.

Если описанная выше логика разделения пакета и разбиения набора данных нежелательна, используйте вместо неё tf.distribute.Strategy.distribute_datasets_from_function, которая не выполняет автоматическое формирование пакетов или разбиение.

Примечание: Если вы используете TPUStrategy, порядок обработки данных рабочими процессами при использовании tf.distribute.Strategy.experimental_distribute_dataset или tf.distribute.Strategy.distribute_datasets_from_function не гарантируется. Это обычно требуется, если вы используете tf.distribute для масштабирования предсказаний. Однако вы можете вставить индекс для каждого элемента в пакете и упорядочить результаты соответствующим образом. Обратитесь к этому фрагменту для примера того, как упорядочить результаты.

Примечание: Состоятельные преобразования наборов данных в настоящее время не поддерживаются с tf.distribute.experimental_distribute_dataset или tf.distribute.distribute_datasets_from_function. Любые состоятельные операции, которые может содержать набор данных, в настоящее время игнорируются. Например, если ваш набор данных имеет map_fn , который использует tf.random.uniform для поворота изображения, то у вас есть граф набора данных, зависящий от состояния (т.е. случайного зерна) на локальной машине, где выполняется процесс Python.

Для получения дополнительной информации о применении и свойствах этого метода обратитесь к тьюториалу по распределённому вводу. Если вы заинтересованы в обработке последних частичных пакетов, прочтите эту секцию.

experimental_local_results

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

experimental_local_results(
    value
)

Возвращает список всех локальных значений на реплику, содержащихся в value.

Примечание: Это возвращает только значения на рабочем процессе, инициированном этим клиентом. При использовании tf.distribute.Strategy, такого как tf.distribute.experimental.MultiWorkerMirroredStrategy, каждый рабочий процесс будет своим клиентом, и эта функция вернёт только значения, вычисленные на этом рабочем процессе.

experimental_make_numpy_dataset

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

experimental_make_numpy_dataset(
    numpy_input, session=None
)

Создаёт tf.data.Dataset для ввода, предоставленного через массив NumPy.

Это позволяет избежать добавления numpy_input в качестве большой константы в графе и копирует данные на машину или машины, которые будут обрабатывать вход.

Обратите внимание, что вам, скорее всего, понадобится использовать tf.distribute.Strategy.experimental_distribute_dataset с возвращаемым набором данных для дальнейшего распределения его с помощью стратегии.

Пример:

numpy_input = np.ones([10], dtype=np.float32)
dataset = strategy.experimental_make_numpy_dataset(numpy_input)
dist_dataset = strategy.experimental_distribute_dataset(dataset)

experimental_run

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

experimental_run(
    fn, input_iterator=None
)

Выполняет операции в fn на каждой реплике, с данными из input_iterator. (устарело)

При включённом режиме выполнения операций, выполняет операции, указанные в fn на каждой реплике. В противном случае создаёт граф для выполнения операций на каждой реплике.

Каждая реплика будет использовать один, отличающийся вход из входов, предоставленных одним вызовом get_next итератора ввода.

fn может вызвать tf.distribute.get_replica_context() для доступа к элементам, таким как replica_id_in_sync_group.

make_dataset_iterator

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

make_dataset_iterator(
    dataset
)

Создаёт итератор для входных данных, предоставленных через dataset.

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

Пользователь также может использовать make_input_fn_iterator, если хочет настроить, какие входные данные будут передаваться каждой реплике/рабочему узлу и т. д.

make_input_fn_iterator

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

make_input_fn_iterator(
    input_fn,
    replication_mode=tf.distribute.InputReplicationMode.PER_WORKER
)

Возвращает итератор, разделенный по репликам, созданный из функции ввода.

Функция input_fn должна принимать объект tf.distribute.InputContext, где можно получить информацию о группировании по пакетам и фрагментации ввода:

def input_fn(input_context):
  batch_size = input_context.get_per_replica_batch_size(global_batch_size)
  d = tf.data.Dataset.from_tensors([[1.]]).repeat().batch(batch_size)
  return d.shard(input_context.num_input_pipelines,
                 input_context.input_pipeline_id)
with strategy.scope():
  iterator = strategy.make_input_fn_iterator(input_fn)
  replica_results = strategy.experimental_run(replica_fn, iterator)

Набор данных tf.data.Dataset, возвращаемый input_fn, должен иметь размер пакета для каждой реплики, который можно вычислить с помощью input_context.get_per_replica_batch_size.

reduce

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

reduce(
    reduce_op, value, axis=None
)

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

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

Учитывая значение для каждой реплики, возвращаемое 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>
}

1. Входной DistributedValues.

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

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

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

scope

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

scope()

Блок кода для установки стратегии в текущий контекст и распределения переменных.

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

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

Что происходит при входе в Strategy.scope?

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

Примечание: Вход в блок не автоматически распределяет вычисления, за исключением случаев использования высокоуровневой обучающей среды, например, Keras model.fit. Если вы не используете model.fit, вам необходимо использовать API strategy.run для явного распределения вычислений. См. пример в руководстве по созданию пользовательского цикла обучения https://www.tensorflow.org/tutorials/distribute/custom_training.

Что должно находиться внутри блока и что вне?

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

• Всё, что создаёт переменные, которые должны быть распределёнными переменными, должно вызываться в пределах strategy.scope. Это можно сделать, либо вызвав функцию создания переменной непосредственно внутри контекста блока, либо, полагаясь на другой API, например, strategy.run или 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 https://www.tensorflow.org/tutorials/distribute/keras. ВНИМАНИЕ: Просто вызов model(..) не автоматически вводит пойманный блок — только API высокоуровневой обучающей среды поддерживают это поведение: model.compile, model.fit, model.evaluate, model.predict и model.save могут вызываться как внутри, так и вне блока.
• Следующее может быть как внутри, так и вне блока:
  • Создание наборов данных ввода
  • Определение tf.functions, представляющих шаг обучения
  • API сохранения, такие как tf.saved_model.save. Загрузка создаёт переменные, поэтому это должно быть внутри блока, если вы хотите обучить модель распределённым способом.
  • Сохранение контрольных точек. Как упоминалось выше — checkpoint.restore иногда может потребоваться внутри блока, если оно создаёт переменные.

update_config_proto

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

update_config_proto(
    config_proto
)

Возвращает копию config_proto с изменениями для использования с этой стратегией.

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