Класс MirroredStrategy

Класс MirroredStrategy представляет собой стратегию распределения, которая синхронно выполняет вычисления на нескольких GPU, расположенных на одной машине. Каждое устройство получает копию модели и переменных, а градиенты усредняются между всеми устройствами после каждого шага. Конструктор принимает список устройств, параметры кросс-девайс операций и дополнительные опции. Если список устройств не указан, стратегия автоматически обнаружит все доступные GPU.

Основное применение класса - обёртка переменных модели и датасетов для обучения на нескольких GPU. Переменные создаются внутри метода scope, а данные распределяются через experimental_distribute_dataset или distribute_datasets_from_function.

Синтаксис

tf.distribute.MirroredStrategy(
	devices=None,
	cross_device_ops=None,
	cluster_spec=None,
	compute_devices=None,
	communication_options=None
)

Пример

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

import tensorflow as tf

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

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

1

Пример

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

import tensorflow as tf

tf.random.set_seed(0)

strategy = tf.distribute.MirroredStrategy()

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

dataset = tf.data.Dataset.from_tensor_slices(
	(tf.constant([[1.0], [2.0], [3.0], [4.0]]),
	 tf.constant([[2.0], [4.0], [6.0], [8.0]]))
).batch(2)

dist_dataset = strategy.experimental_distribute_dataset(dataset)

model.fit(dist_dataset, epochs=1)

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

1/1 [==============================] - 0s 200ms/step - loss: 15.2345

Пример

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

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
	t = tf.constant([1, 2, 3, 4, 5])
	res = strategy.run(lambda x: x + 1, args=(t,))

print(res)

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

tf.Tensor([2 3 4 5 6], shape=(5,), dtype=int32)

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

метод scope,
который создает область видимости для распределённых переменных
метод run,
который выполняет функцию на каждом устройстве стратегии
метод reduce,
который объединяет результаты вычислений с устройств
метод experimental_distribute_dataset,
который распределяет датасет между устройствами