Метод run класса MirroredStrategy

Метод run класса MirroredStrategy выполняет указанную функцию на каждой реплике (устройстве) распределенной стратегии. Первым параметром метод принимает функцию, которую нужно выполнить, а последующими параметрами - аргументы этой функции. Метод возвращает результат выполнения функции на каждой из реплик в виде списка.

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

Синтаксис

strategy.run(fn, args=(), kwargs=None)

Пример

Давайте создадим стратегию и выполним простую функцию на всех репликах:

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()

def fn():
	return tf.constant([1, 2, 3, 4, 5])

res = strategy.run(fn)
print(res)

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

PerReplica:{
  0: tf.Tensor([1 2 3 4 5], shape=(5,), dtype=int32)
}

Пример

Давайте передадим аргументы в функцию, выполняемую на репликах:

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()

def fn(x):
	return x * 2

t = tf.constant([1, 2, 3, 4, 5])
res = strategy.run(fn, args=(t,))
print(res)

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

PerReplica:{
  0: tf.Tensor([ 2  4  6  8 10], shape=(5,), dtype=int32)
}

Пример

Давайте выполним функцию с несколькими аргументами и именованными параметрами:

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()

def fn(a, b, c=1):
	return a + b + c

t1 = tf.constant([1, 2, 3, 4, 5])
t2 = tf.constant([10, 20, 30, 40, 50])
res = strategy.run(fn, args=(t1, t2), kwargs={'c': 100})
print(res)

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

PerReplica:{
  0: tf.Tensor([111 122 133 144 155], shape=(5,), dtype=int32)
}

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

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