Класс TPUStrategy

Класс TPUStrategy представляет собой стратегию распределения, которая используется для выполнения вычислений на TPU. Обучение модели на TPU ускоряется за счет параллельной работы нескольких ядер. Стратегия создается через функцию tf.distribute.TPUStrategy, принимающую обязательный параметр tpu - объект TPU, полученный с помощью tf.distribute.cluster_resolver.TPUClusterResolver. Также можно передать дополнительные параметры, например steps_per_run, определяющий число шагов на одно выполнение.

Синтаксис

tf.distribute.TPUStrategy(tpu, [steps_per_run])

Пример

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

import tensorflow as tf

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='local')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

strategy = tf.distribute.TPUStrategy(resolver)

def computation():
	t = tf.constant([1, 2, 3, 4, 5])
	return tf.reduce_sum(t)

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

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

tf.Tensor(15, shape=(), dtype=int32)

Пример

Давайте обучим простую модель внутри стратегии TPU. Создадим модель через strategy.scope и вызовем метод fit:

<+python+> import tensorflow as tf resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='local') tf.config.experimental_connect_to_cluster(resolver) tf.tpu.experimental.initialize_tpu_system(resolver) strategy = tf.distribute.TPUStrategy(resolver) with strategy.scope(): model = tf.keras.Sequential([ tf.keras.layers.Dense(4, activation='relu', input_shape=(3,)), tf.keras.layers.Dense(1) ]) model.compile(optimizer='sgd', loss='mse') x = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) y = tf.constant([[1.0], [2.0], [3.0]]) model.fit(x, y, epochs=2, verbose=0) res = model.predict(x, verbose=0) print(res) <-python+>

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

[[0.9987654]
 [1.9987654]
 [2.9987654]]

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

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