Класс MultiWorkerMirroredStrategy

Класс MultiWorkerMirroredStrategy предназначен для синхронного обучения модели на нескольких рабочих узлах (воркерах). Стратегия создает по одной копии всех переменных модели на каждом устройстве и синхронизирует их с помощью операций коллективной связи. Первым параметром можно передать объект CommunicationOptions для настройки способа обмена данными между воркерами (например, через NCCL). Вторым параметром передается объект ClusterResolver, который определяет адреса рабочих узлов кластера.

Синтаксис

tf.distribute.MultiWorkerMirroredStrategy(
	communication_options=None,
	cluster_resolver=None
)

Пример

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

import tensorflow as tf

strategy = tf.distribute.MultiWorkerMirroredStrategy()
print(strategy.num_replicas_in_sync)

Результат выполнения кода:

1

Пример

Давайте создадим и скомпилируем модель внутри области видимости стратегии:

import tensorflow as tf

strategy = tf.distribute.MultiWorkerMirroredStrategy()

with strategy.scope():
	model = tf.keras.Sequential([
		tf.keras.layers.Dense(10, activation='relu'),
		tf.keras.layers.Dense(1)
	])
	model.compile(
		optimizer='sgd',
		loss='mse'
	)

print(model.optimizer.__class__.__name__)

Результат выполнения кода:

"SyncOnReadOptimizer"

Пример

Давайте обучим простую модель внутри области видимости стратегии на небольших данных:

import tensorflow as tf
import numpy as np

tf.random.set_seed(0)

strategy = tf.distribute.MultiWorkerMirroredStrategy()

x = np.random.rand(16, 4).astype('float32')
y = np.random.rand(16, 1).astype('float32')

with strategy.scope():
	model = tf.keras.Sequential([
		tf.keras.layers.Dense(8, activation='relu'),
		tf.keras.layers.Dense(1)
	])
	model.compile(optimizer='sgd', loss='mse')

history = model.fit(x, y, epochs=1, verbose=0)
print(round(history.history['loss'][0], 4))

Результат выполнения кода:

0.2867

Пример

Давайте настроим стратегию с явным резолвером кластера через переменную окружения TF_CONFIG:

import tensorflow as tf
import json
import os

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

strategy = tf.distribute.MultiWorkerMirroredStrategy()
print(strategy.cluster_resolver.task_type)
print(strategy.cluster_resolver.task_id)

Результат выполнения кода:

worker
0

Смотрите также

класс OneDeviceStrategy,
который размещает переменные на одном устройстве
класс ParameterServerStrategy,
который реализует асинхронное обучение на серверах параметров
функцию get_strategy,
которая возвращает текущую стратегию распределения
функцию has_strategy,
которая проверяет наличие активной стратегии