tf.distribute.experimental.ParameterServerStrategy

Стратегия tf.distribute для нескольких рабочих узлов с серверами параметров.

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

tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver: tf.distribute.cluster_resolver.ClusterResolver,
    variable_partitioner: tf.distribute.experimental.partitioners.Partitioner = None
)

Используется в блокнотах

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

В TensorFlow 2 для обучения с серверами параметров рекомендуется архитектура с централизованной координацией. На каждом рабочем узле и сервере параметров запускается tf.distribute.Server, а задача координатора отвечает за создание ресурсов на рабочих узлах и серверах параметров, отправку функций и координацию обучения. Координатор использует tf.distribute.experimental.coordinator.ClusterCoordinator для координации кластера и tf.distribute.experimental.ParameterServerStrategy для определения переменных на серверах параметров и вычислений на рабочих узлах.

Для выполнения обучения координатор отправляет tf.functionы на удаленные рабочие узлы. Получив запросы от координатора, рабочий узел выполняет tf.function, считывая переменные с серверов параметров, выполняя операции и обновляя переменные на серверах параметров. Каждый рабочий узел обрабатывает только запросы координатора и взаимодействует с серверами параметров, не взаимодействуя напрямую с другими рабочими узлами кластера.

В результате сбои некоторых рабочих узлов не мешают кластеру продолжать работу, что позволяет обучать модель на экземплярах, которые иногда могут быть недоступны (например, прерываемых или спотовых экземплярах). Однако координатор и серверы параметров должны быть доступны постоянно, чтобы кластер мог продолжать работу.

Обратите внимание, что координатор не является одним из рабочих узлов обучения. Вместо этого он создает такие ресурсы, как переменные и наборы данных, отправляет tf.functionы, сохраняет контрольные точки и т. д. Помимо рабочих узлов, серверов параметров и координатора, можно дополнительно запустить оценщик, который периодически считывает сохраненные координатором контрольные точки и выполняет оценку каждой из них.

ParameterServerStrategy поддерживает два API для обучения: пользовательский цикл обучения (CTL) и API обучения Keras, также известный как Model.fit. CTL рекомендуется использовать, если пользователи предпочитают самостоятельно определять детали цикла обучения, а Model.fit рекомендуется, если предпочтительны высокоуровневая абстракция и автоматическая обработка обучения.

При использовании CTL ParameterServerStrategy необходимо применять вместе с объектом tf.distribute.experimental.coordinator.ClusterCoordinator.

При использовании Model.fit в настоящее время поддерживается только тип входных данных tf.keras.utils.experimental.DatasetCreator.

Пример кода для координатора

В этом разделе приведены фрагменты кода, предназначенные для запуска на единственной задаче, назначенной координатором. Обратите внимание, что аргументы cluster_resolver, variable_partitioner и dataset_fn описаны в следующих разделах: «Настройка кластера», «Разбиение переменных на разделы» и «Подготовка набора данных».

С CTL:

# Prepare a strategy to use with the cluster and variable partitioning info.
strategy = tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver=...,
    variable_partitioner=...)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy=strategy)

# Prepare a distribute dataset that will place datasets on the workers.
distributed_dataset = coordinator.create_per_worker_dataset(dataset_fn=...)

with strategy.scope():
  model = ...
  optimizer, metrics = ...  # Keras optimizer/metrics are great choices
  checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
  checkpoint_manager = tf.train.CheckpointManager(
      checkpoint, checkpoint_dir, max_to_keep=2)
  # `load_checkpoint` infers initial epoch from `optimizer.iterations`.
  initial_epoch = load_checkpoint(checkpoint_manager) or 0

@tf.function
def worker_fn(iterator):

  def replica_fn(inputs):
    batch_data, labels = inputs
    # calculate gradient, applying gradient, metrics update etc.

  strategy.run(replica_fn, args=(next(iterator),))

for epoch in range(initial_epoch, num_epoch):
  distributed_iterator = iter(distributed_dataset)  # Reset iterator state.
  for step in range(steps_per_epoch):

    # Asynchronously schedule the `worker_fn` to be executed on an arbitrary
    # worker. This call returns immediately.
    coordinator.schedule(worker_fn, args=(distributed_iterator,))

  # `join` blocks until all scheduled `worker_fn`s finish execution. Once it
  # returns, we can read the metrics and save checkpoints as needed.
  coordinator.join()
  logging.info('Metric result: %r', metrics.result())
  train_accuracy.reset_states()
  checkpoint_manager.save()

С Model.fit:

# Prepare a strategy to use with the cluster and variable partitioning info.
strategy = tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver=...,
    variable_partitioner=...)

# A dataset function takes a `input_context` and returns a `Dataset`
def dataset_fn(input_context):
  dataset = tf.data.Dataset.from_tensors(...)
  return dataset.repeat().shard(...).batch(...).prefetch(...)

# With `Model.fit`, a `DatasetCreator` needs to be used.
input = tf.keras.utils.experimental.DatasetCreator(dataset_fn=...)

with strategy.scope():
  model = ...  # Make sure the `Model` is created within scope.
model.compile(optimizer="rmsprop", loss="mse", steps_per_execution=..., ...)

# Optional callbacks to checkpoint the model, back up the progress, etc.
callbacks = [tf.keras.callbacks.ModelCheckpoint(...), ...]

# `steps_per_epoch` is required with `ParameterServerStrategy`.
model.fit(input, epochs=..., steps_per_epoch=..., callbacks=callbacks)

Пример кода для рабочего узла и серверов параметров

Помимо координатора, необходимо назначить задачи типа «worker» или «ps». Для запуска сервера TensorFlow, ожидающего запросов координатора, они должны выполнить следующий код:

# Provide a `tf.distribute.cluster_resolver.ClusterResolver` that serves
# the cluster information. See below "Cluster setup" section.
cluster_resolver = ...

server = tf.distribute.Server(
    cluster_resolver.cluster_spec(),
    job_name=cluster_resolver.task_type,
    task_index=cluster_resolver.task_id,
    protocol="grpc")

# Blocking the process that starts a server from exiting.
server.join()

Настройка кластера

Чтобы задачи кластера знали адреса других задач, в координаторе, рабочем узле и на сервере параметров необходимо использовать tf.distribute.cluster_resolver.ClusterResolver. Объект tf.distribute.cluster_resolver.ClusterResolver предоставляет сведения о кластере, а также тип и идентификатор текущей задачи. Дополнительную информацию см. в разделе tf.distribute.cluster_resolver.ClusterResolver.

Если задана переменная среды TF_CONFIG, также следует использовать tf.distribute.cluster_resolver.TFConfigClusterResolver.

Поскольку в tf.distribute.experimental.ParameterServerStrategy предполагаются определенные имена типов задач, в tf.distribute.cluster_resolver.ClusterResolver для обозначения координатора, серверов параметров и рабочих узлов следует использовать имена «chief», «ps» и «worker» соответственно.

В следующем примере задается TF_CONFIG для задачи, назначенной сервером параметров (тип задачи «ps») с индексом 1 (вторая задача) в кластере с одним координатором, двумя серверами параметров и тремя рабочими узлами. Обратите внимание: ее необходимо задать до использования tf.distribute.cluster_resolver.TFConfigClusterResolver.

Пример кода настройки кластера:

os.environ['TF_CONFIG'] = '''
{
  "cluster": {
    "chief": ["chief.example.com:2222"],
    "ps": ["ps0.example.com:2222", "ps1.example.com:2222"],
    "worker": ["worker0.example.com:2222", "worker1.example.com:2222",
               "worker2.example.com:2222"]
  },
  "task": {
    "type": "ps",
    "index": 1
  }
}
'''

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

# If coordinator, create a strategy and start the training program.
if cluster_resolver.task_type == 'chief':
  strategy = tf.distribute.experimental.ParameterServerStrategy(
      cluster_resolver)
  ...

# If worker/ps, create a server
elif cluster_resolver.task_type in ("worker", "ps"):
  server = tf.distribute.Server(...)
  ...

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

Создание переменных с помощью strategy.scope()

tf.distribute.experimental.ParameterServerStrategy соответствует контракту API tf.distribute, согласно которому переменные необходимо создавать внутри контекстного менеджера, возвращаемого strategy.scope(), чтобы они правильно размещались на серверах параметров по круговому принципу:

# In this example, we're assuming having 3 ps.
strategy = tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver=...)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy=strategy)

# Variables should be created inside scope to be placed on parameter servers.
# If created outside scope such as `v1` here, it would be placed on the
# coordinator.
v1 = tf.Variable(initial_value=0.0)

with strategy.scope():
  v2 = tf.Variable(initial_value=1.0)
  v3 = tf.Variable(initial_value=2.0)
  v4 = tf.Variable(initial_value=3.0)
  v5 = tf.Variable(initial_value=4.0)

# v2 through v5 are created in scope and are distributed on parameter servers.
# Default placement is round-robin but the order should not be relied on.
assert v2.device == "/job:ps/replica:0/task:0/device:CPU:0"
assert v3.device == "/job:ps/replica:0/task:1/device:CPU:0"
assert v4.device == "/job:ps/replica:0/task:2/device:CPU:0"
assert v5.device == "/job:ps/replica:0/task:0/device:CPU:0"

Дополнительную информацию см. в разделе distribute.Strategy.scope.

Разбиение переменных на разделы

Наличие выделенных серверов для хранения переменных позволяет разделять их, или «шардировать», между серверами ps. Разбиение больших переменных между серверами ps — широко используемый метод повышения пропускной способности обучения и снижения ограничений по памяти. Он позволяет параллельно вычислять и обновлять разные шарды переменной и зачастую обеспечивает более равномерное распределение нагрузки между серверами параметров. Без шардирования невозможно обучать модели с большими переменными (например, эмбеддингами), которые не помещаются в память одной машины.

В tf.distribute.experimental.ParameterServerStrategy, если для __init__ предоставлен variable_partitioner и выполнены определенные условия, созданные в области действия переменные распределяются по серверам параметров по круговому принципу. Ссылка на переменную, возвращаемая tf.Variable, становится типом, служащим контейнером для шардированных переменных. Для доступа к фактическим компонентам переменной можно использовать атрибут variables этого контейнера. При создании модели с помощью tf.Module или Keras компоненты переменной собираются в атрибутах, аналогичных variables.

Рекомендуется использовать разделители, основанные на размере, например tf.distribute.experimental.partitioners.MinSizePartitioner, чтобы избежать разбиения небольших переменных, которое может негативно повлиять на скорость обучения модели.

# Partition the embedding layer into 2 shards.
variable_partitioner = (
  tf.distribute.experimental.partitioners.MinSizePartitioner(
    min_shard_bytes=(256 << 10),
    max_shards = 2))
strategy = tf.distribute.experimental.ParameterServerStrategy(
  cluster_resolver=...,
  variable_partitioner = variable_partitioner)
with strategy.scope():
  embedding = tf.keras.layers.Embedding(input_dim=1024, output_dim=1024)
assert len(embedding.variables) == 2
assert isinstance(embedding.variables[0], tf.Variable)
assert isinstance(embedding.variables[1], tf.Variable)
assert embedding.variables[0].shape == (512, 1024)
assert embedding.variables[1].shape == (512, 1024)

Контейнер шардированной переменной можно преобразовать в Tensor с помощью tf.convert_to_tensor. Это означает, что контейнер можно напрямую использовать в большинстве операций Python, где такое преобразование в Tensor выполняется автоматически. Например, в приведенном выше фрагменте кода x * self.w неявно выполнит это преобразование в тензор. Обратите внимание, что такое преобразование может быть затратным, поскольку компоненты переменной необходимо передать с нескольких серверов параметров туда, где используется значение.

С другой стороны, tf.nn.embedding_lookup не выполняет преобразование в тензор, а вместо этого параллельно ищет значения в компонентах переменной. Это крайне важно для масштабирования операций поиска в эмбеддингах, когда таблица эмбеддингов содержит большую переменную.

При сохранении разделенной на части переменной в SavedModel она сохраняется как единая переменная. Это повышает эффективность обслуживания, устраняя ряд операций, обрабатывающих особенности разбиения на части.

Известные ограничения разбиения переменных на части:

• Количество разделов не должно меняться между сохранением и загрузкой контрольной точки.

• После сохранения разделенных на части переменных в SavedModel его нельзя загрузить с помощью tf.saved_model.load.

• Разбиение переменных на части не работает напрямую с tf.GradientTape. Используйте атрибуты variables, чтобы получить фактические компоненты переменной и передать их в API для работы с градиентами.

Подготовка набора данных

При использовании tf.distribute.experimental.ParameterServerStrategy на каждом рабочем узле создается набор данных для обучения. Для этого создается dataset_fn, которая не принимает аргументов и возвращает tf.data.Dataset, а затем передается dataset_fn в tf.distribute.experimental.coordinator. ClusterCoordinator.create_per_worker_dataset. Рекомендуется перемешивать и повторять набор данных, чтобы обучающая выборка обрабатывалась как можно равномернее.

def dataset_fn():
  filenames = ...
  dataset = tf.data.Dataset.from_tensor_slices(filenames)

  # Dataset is recommended to be shuffled, and repeated.
  return dataset.shuffle(buffer_size=...).repeat().batch(batch_size=...)

coordinator =
    tf.distribute.experimental.coordinator.ClusterCoordinator(strategy=...)
distributed_dataset = coordinator.create_per_worker_dataset(dataset_fn)

Ограничения

• tf.distribute.experimental.ParameterServerStrategy в TF2 является экспериментальной, и ее API может измениться.

• При использовании Model.fit стратегию tf.distribute.experimental.ParameterServerStrategy необходимо использовать с tf.keras.utils.experimental.DatasetCreator, а параметр steps_per_epoch должен быть указан.

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

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

...

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

Методы

distribute_datasets_from_function

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

distribute_datasets_from_function(
    dataset_fn, options=None
)

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

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

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

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

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

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

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

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

experimental_distribute_dataset

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

experimental_distribute_dataset(
    dataset, options=None
)

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

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

Пример:

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

Этот метод выполняет три ключевые операции: разбиение на пакеты, разбиение на шарды и предварительную выборку.

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

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

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

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

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

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

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

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

experimental_distribute_values_from_function

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

experimental_distribute_values_from_function(
    value_fn
)

Создает tf.distribute.DistributedValues из value_fn.

Эта функция предназначена для создания tf.distribute.DistributedValues, передаваемых в run, reduce или другие методы, принимающие распределенные значения, когда наборы данных не используются.

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

1. Возврат постоянного значения для каждой реплики:

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   def value_fn(ctx):
     return tf.constant(1.)
   distributed_values = (
       strategy.experimental_distribute_values_from_function(
          value_fn))
   local_result = strategy.experimental_local_results(
       distributed_values)
   local_result
       (<tf.Tensor: shape=(), dtype=float32, numpy=1.0>,
       <tf.Tensor: shape=(), dtype=float32, numpy=1.0>)
       

2. Распределение значений массива на основе replica_id:

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   array_value = np.array([3., 2., 1.])
   def value_fn(ctx):
     return array_value[ctx.replica_id_in_sync_group]
   distributed_values = (
       strategy.experimental_distribute_values_from_function(
           value_fn))
   local_result = strategy.experimental_local_results(
       distributed_values)
   local_result
       (3.0, 2.0)
       

3. Указание значений с помощью num_replicas_in_sync:

   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   def value_fn(ctx):
     return ctx.num_replicas_in_sync
   distributed_values = (
       strategy.experimental_distribute_values_from_function(
           value_fn))
   local_result = strategy.experimental_local_results(
       distributed_values)
   local_result
       (2, 2)
       

4. Размещение значений на устройствах и их распределение:

   strategy = tf.distribute.TPUStrategy()
   worker_devices = strategy.extended.worker_devices
   multiple_values = []
   for i in range(strategy.num_replicas_in_sync):
     with tf.device(worker_devices[i]):
       multiple_values.append(tf.constant(1.0))
   
   def value_fn(ctx):
     return multiple_values[ctx.replica_id_in_sync_group]
   
   distributed_values = strategy.
     experimental_distribute_values_from_function(
     value_fn)
   

experimental_local_results

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

experimental_local_results(
    value
)

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

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

gather

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

gather(
    value, axis
)

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

Для объекта типа tf.distribute.DistributedValues или tf.Tensor value этот API собирает и объединяет конкатенацией value по репликам вдоль измерения axis. Результат копируется на «текущее» устройство, которым обычно является ЦП рабочего узла, на котором выполняется программа. Для 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 и двумя 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 это ЦП каждого рабочего узла.

В tf.distribute есть несколько API для редукции значений по репликам:

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

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

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

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

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

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

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

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

run

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

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

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

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

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

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

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

1. Входной тензор-константа.

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

2. Входные данные DistributedValues.

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

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

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

scope

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

scope()

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

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

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

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

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

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

Что следует размещать внутри области видимости, а что — за ее пределами?

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

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