¿Tengo una pregunta? Conéctese con la comunidad en el Foro de visita del foro de TensorFlow

Usar TPU

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar cuaderno

El soporte experimental para Cloud TPU está disponible actualmente para Keras y Google Colab. Antes de ejecutar este portátil Colab, asegúrese de que su acelerador de hardware sea una TPU comprobando la configuración de su portátil: Tiempo de ejecución > Cambiar tipo de tiempo de ejecución > Acelerador de hardware > TPU .

Configuración

import tensorflow as tf

import os
import tensorflow_datasets as tfds

Inicialización de TPU

Las TPU suelen ser trabajadores de Cloud TPU, que son diferentes del proceso local que ejecuta el programa Python del usuario. Por lo tanto, debe realizar un trabajo de inicialización para conectarse al clúster remoto e inicializar las TPU. Tenga en cuenta que el argumento tpu para tf.distribute.cluster_resolver.TPUClusterResolver es una dirección especial solo para Colab. Si está ejecutando su código en Google Compute Engine (GCE), debe pasar el nombre de su Cloud TPU.

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.26:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.26:8470
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
All devices:  [LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:5', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:3', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:0', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:1', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:2', device_type='TPU')]

Colocación manual del dispositivo

Una vez que se inicializa la TPU, puede usar la ubicación manual del dispositivo para colocar el cálculo en un solo dispositivo de TPU:

a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

with tf.device('/TPU:0'):
  c = tf.matmul(a, b)

print("c device: ", c.device)
print(c)
c device:  /job:worker/replica:0/task:0/device:TPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

Estrategias de distribución

Por lo general, ejecuta su modelo en varias TPU de forma paralela a los datos. Para distribuir su modelo en varias TPU (u otros aceleradores), TensorFlow ofrece varias estrategias de distribución. Puede reemplazar su estrategia de distribución y el modelo se ejecutará en cualquier dispositivo (TPU) dado. Consulte la guía de estrategia de distribución para obtener más información.

Para demostrar esto, cree un objeto tf.distribute.TPUStrategy :

strategy = tf.distribute.TPUStrategy(resolver)
INFO:tensorflow:Found TPU system:
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

Para replicar un cálculo para que se pueda ejecutar en todos los núcleos de TPU, puede pasarlo a la API de strategy.run . A continuación se muestra un ejemplo que muestra todos los núcleos que reciben las mismas entradas (a, b) y realizan la multiplicación de matrices en cada núcleo de forma independiente. Las salidas serán los valores de todas las réplicas.

@tf.function
def matmul_fn(x, y):
  z = tf.matmul(x, y)
  return z

z = strategy.run(matmul_fn, args=(a, b))
print(z)
PerReplica:{
  0: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  1: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  2: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  3: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  4: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  5: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  6: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  7: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)
}

Clasificación en TPU

Habiendo cubierto los conceptos básicos, considere un ejemplo más concreto. En esta sección, se muestra cómo usar la estrategia de distribución, tf.distribute.TPUStrategy para entrenar un modelo de Keras en una Cloud TPU.

Definir un modelo de Keras

Comience con una definición de un modelo Sequential Keras para la clasificación de imágenes en el conjunto de datos MNIST usando Keras. No es diferente de lo que usaría si estuviera entrenando en CPU o GPU. Tenga en cuenta que la creación del modelo de Keras debe estar dentro de strategy.scope , por lo que las variables se pueden crear en cada dispositivo de TPU. No es necesario que otras partes del código estén dentro del alcance de la estrategia.

def create_model():
  return tf.keras.Sequential(
      [tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)),
       tf.keras.layers.Conv2D(256, 3, activation='relu'),
       tf.keras.layers.Flatten(),
       tf.keras.layers.Dense(256, activation='relu'),
       tf.keras.layers.Dense(128, activation='relu'),
       tf.keras.layers.Dense(10)])

Cargar el conjunto de datos

El uso eficiente de la APItf.data.Dataset es fundamental cuando se usa una TPU de Cloud, ya que es imposible usar las TPU de Cloud a menos que puedas proporcionarles datos con la suficiente rapidez. Puede obtener más información sobre el rendimiento del conjunto de datos en la guía de rendimiento de la canalización de entrada .

Para todos los experimentos, excepto los más simples (con tf.data.Dataset.from_tensor_slices u otros datos en el gráfico), debes almacenar todos los archivos de datos leídos por el conjunto de datos en depósitos de Google Cloud Storage (GCS).

Para la mayoría de los casos de uso, se recomienda convertir sus datos al formato TFRecord y usar un tf.data.TFRecordDataset para leerlos. Consulte el tutorial TFRecord y tf.Example para obtener detalles sobre cómo hacer esto. No es un requisito estricto y puede utilizar otros lectores de conjuntos de datos, como tf.data.FixedLengthRecordDataset o tf.data.TextLineDataset .

Puede cargar pequeños conjuntos de datos completos en la memoria usando tf.data.Dataset.cache .

Independientemente del formato de datos utilizado, se recomienda encarecidamente que utilice archivos grandes del orden de 100 MB. Esto es especialmente importante en esta configuración de red, ya que la sobrecarga de abrir un archivo es significativamente mayor.

Como se muestra en el código a continuación, debe usar el módulo tensorflow_datasets para obtener una copia de los datos de prueba y entrenamiento de MNIST. Tenga en cuenta que try_gcs se especifica para usar una copia que está disponible en un depósito de GCS público. Si no especifica esto, la TPU no podrá acceder a los datos descargados.

def get_dataset(batch_size, is_training=True):
  split = 'train' if is_training else 'test'
  dataset, info = tfds.load(name='mnist', split=split, with_info=True,
                            as_supervised=True, try_gcs=True)

  # Normalize the input data.
  def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255.0
    return image, label

  dataset = dataset.map(scale)

  # Only shuffle and repeat the dataset in training. The advantage of having an
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so that you don't need to think about scaling the gradients
  # based on the actual batch size.
  if is_training:
    dataset = dataset.shuffle(10000)
    dataset = dataset.repeat()

  dataset = dataset.batch(batch_size)

  return dataset

Entrene el modelo con las API de alto nivel de Keras

Puede entrenar su modelo con las API de compile y fit Keras. No hay nada específico de TPU en este paso: escribe el código como si estuviera usando varias GPU y una MirroredStrategy lugar de la TPUStrategy . Puede obtener más información en el tutorial Entrenamiento distribuido con Keras .

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

batch_size = 200
steps_per_epoch = 60000 // batch_size
validation_steps = 10000 // batch_size

train_dataset = get_dataset(batch_size, is_training=True)
test_dataset = get_dataset(batch_size, is_training=False)

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset, 
          validation_steps=validation_steps)
Epoch 1/5
300/300 [==============================] - 17s 32ms/step - loss: 0.1395 - sparse_categorical_accuracy: 0.9570 - val_loss: 0.0564 - val_sparse_categorical_accuracy: 0.9818
Epoch 2/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0331 - sparse_categorical_accuracy: 0.9894 - val_loss: 0.0394 - val_sparse_categorical_accuracy: 0.9877
Epoch 3/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0192 - sparse_categorical_accuracy: 0.9938 - val_loss: 0.0426 - val_sparse_categorical_accuracy: 0.9877
Epoch 4/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0131 - sparse_categorical_accuracy: 0.9956 - val_loss: 0.0442 - val_sparse_categorical_accuracy: 0.9885
Epoch 5/5
300/300 [==============================] - 6s 22ms/step - loss: 0.0098 - sparse_categorical_accuracy: 0.9966 - val_loss: 0.0401 - val_sparse_categorical_accuracy: 0.9879
<tensorflow.python.keras.callbacks.History at 0x7f39a81b05c0>

Para reducir la sobrecarga de Python y maximizar el rendimiento de su TPU, pase el argumento steps_per_execution a Model.compile . En este ejemplo, aumenta el rendimiento en aproximadamente un 50%:

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                # Anything between 2 and `steps_per_epoch` could help here.
                steps_per_execution = 50,
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset,
          validation_steps=validation_steps)
Epoch 1/5
300/300 [==============================] - 12s 40ms/step - loss: 0.1401 - sparse_categorical_accuracy: 0.9566 - val_loss: 0.0524 - val_sparse_categorical_accuracy: 0.9827
Epoch 2/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0331 - sparse_categorical_accuracy: 0.9898 - val_loss: 0.0412 - val_sparse_categorical_accuracy: 0.9873
Epoch 3/5
300/300 [==============================] - 3s 9ms/step - loss: 0.0195 - sparse_categorical_accuracy: 0.9934 - val_loss: 0.0443 - val_sparse_categorical_accuracy: 0.9881
Epoch 4/5
300/300 [==============================] - 3s 9ms/step - loss: 0.0120 - sparse_categorical_accuracy: 0.9962 - val_loss: 0.0433 - val_sparse_categorical_accuracy: 0.9876
Epoch 5/5
300/300 [==============================] - 3s 9ms/step - loss: 0.0104 - sparse_categorical_accuracy: 0.9965 - val_loss: 0.0469 - val_sparse_categorical_accuracy: 0.9874
<tensorflow.python.keras.callbacks.History at 0x7f39a81c5a58>

Entrena el modelo usando un ciclo de entrenamiento personalizado

También puede crear y entrenar su modelo usando las API tf.function y tf.distribute directamente. Puede usar la API strategy.experimental_distribute_datasets_from_function para distribuir el conjunto de datos dada una función de conjunto de datos. Tenga en cuenta que en el ejemplo siguiente, el tamaño del lote que se pasa al conjunto de datos es el tamaño del lote por réplica en lugar del tamaño del lote global. Para obtener más información, consulte el tutorial Capacitación personalizada con tf.distribute.Strategy .

Primero, cree el modelo, los conjuntos de datos y las funciones tf .:

# Create the model, optimizer and metrics inside the strategy scope, so that the
# variables can be mirrored on each device.
with strategy.scope():
  model = create_model()
  optimizer = tf.keras.optimizers.Adam()
  training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)
  training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      'training_accuracy', dtype=tf.float32)

# Calculate per replica batch size, and distribute the datasets on each TPU
# worker.
per_replica_batch_size = batch_size // strategy.num_replicas_in_sync

train_dataset = strategy.experimental_distribute_datasets_from_function(
    lambda _: get_dataset(per_replica_batch_size, is_training=True))

@tf.function
def train_step(iterator):
  """The step function for one training step."""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  strategy.run(step_fn, args=(next(iterator),))
WARNING:tensorflow:From <ipython-input-1-5625c2a14441>:15: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
rename to distribute_datasets_from_function
WARNING:tensorflow:From <ipython-input-1-5625c2a14441>:15: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
rename to distribute_datasets_from_function

Luego, ejecuta el ciclo de entrenamiento:

steps_per_eval = 10000 // batch_size

train_iterator = iter(train_dataset)
for epoch in range(5):
  print('Epoch: {}/5'.format(epoch))

  for step in range(steps_per_epoch):
    train_step(train_iterator)
  print('Current step: {}, training loss: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))
  training_loss.reset_states()
  training_accuracy.reset_states()
Epoch: 0/5
Current step: 300, training loss: 0.135, accuracy: 95.88%
Epoch: 1/5
Current step: 600, training loss: 0.0342, accuracy: 98.89%
Epoch: 2/5
Current step: 900, training loss: 0.0198, accuracy: 99.38%
Epoch: 3/5
Current step: 1200, training loss: 0.0118, accuracy: 99.63%
Epoch: 4/5
Current step: 1500, training loss: 0.0109, accuracy: 99.63%

Mejora del rendimiento con varios pasos dentro de tf.function

Puede mejorar el rendimiento ejecutando varios pasos dentro de una función tf.function . Esto se logra envolviendo la llamada a strategy.run con un tf.range dentro de tf.function , y AutoGraph lo convertirá en un tf.while_loop en el trabajador de TPU.

A pesar del rendimiento mejorado, este método presenta tf.function desventajas en comparación con ejecutar un solo paso dentro de tf.function . Ejecutar varios pasos en una función tf.function es menos flexible; no puede ejecutar cosas con entusiasmo o código de Python arbitrario dentro de los pasos.

@tf.function
def train_multiple_steps(iterator, steps):
  """The step function for one training step."""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  for _ in tf.range(steps):
    strategy.run(step_fn, args=(next(iterator),))

# Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get 
# retraced if the value changes.
train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))

print('Current step: {}, training loss: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))
Current step: 1800, training loss: 0.0075, accuracy: 99.77%

Próximos pasos