¡Reserva! Google I / O regresa del 18 al 20 de mayo Regístrese ahora
Se usó la API de Cloud Translation para traducir esta página.
Switch to English

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 estos portátiles 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 estar en trabajadores de Cloud TPU que son diferentes del proceso local que ejecuta el programa Python del usuario. Por lo tanto, es necesario realizar un trabajo de inicialización para conectarse al clúster remoto e inicializar las TPU. Tenga en cuenta que el argumento tpu para TPUClusterResolver es una dirección especial solo para Colab. En el caso de que esté ejecutando Google Compute Engine (GCE), debe pasar el nombre de su CloudTPU.

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.74:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.74: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

La mayoría de las veces, los usuarios desean ejecutar el modelo en varias TPU de forma paralela a los datos. Una estrategia de distribución es una abstracción que se puede utilizar para impulsar modelos en CPU, GPU o TPU. Simplemente cambie la estrategia de distribución y el modelo se ejecutará en el dispositivo dado. Consulte la guía de estrategia de distribución para obtener más información.

Primero, crea el objeto 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 de modo que pueda ejecutarse en todos los núcleos de TPU, simplemente puede pasarlo a la API strategy.run . A continuación se muestra un ejemplo de que todos los núcleos obtendrán las mismas entradas (a, b) y realizarán el matmul 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

Como hemos aprendido los conceptos básicos, es hora de mirar un ejemplo más concreto. En esta guía, se muestra cómo usar la estrategia de distribución tf.distribute.TPUStrategy para impulsar una TPU de Cloud y entrenar un modelo de Keras.

Definir un modelo de Keras

A continuación se muestra la definición del modelo MNIST usando Keras, sin cambios de lo que usaría 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)])

Conjuntos de datos de entrada

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. Consulte la Guía de rendimiento de canalización de entrada para obtener detalles sobre el rendimiento del conjunto de datos.

Para toda la experimentación, excepto la más simple (con tf.data.Dataset.from_tensor_slices u otros datos en el gráfico), deberá 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 en formato TFRecord y usar un tf.data.TFRecordDataset para leerlos. Consulte TFRecord y tf.Example tutorial para obtener detalles sobre cómo hacer esto. Sin embargo, esto no es un requisito estricto y puede utilizar otros lectores de conjuntos de datos ( FixedLengthRecordDataset o TextLineDataset ) si lo prefiere.

Los conjuntos de datos pequeños se pueden cargar completamente 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.

Aquí debe usar el módulo tensorflow_datasets para obtener una copia de los datos de 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 que se descargan.

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)

  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 to have a
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so users 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 un modelo usando las API de alto nivel de Keras

Puede entrenar un modelo simplemente con las API de ajuste / compilación de Keras. Nada aquí es específico de TPU, escribiría el mismo código a continuación si tuviera varias GPU y utilizara MirroredStrategy lugar de TPUStrategy . Para obtener más información, consulte 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 [==============================] - 19s 38ms/step - loss: 0.3196 - sparse_categorical_accuracy: 0.8996 - val_loss: 0.0467 - val_sparse_categorical_accuracy: 0.9855
Epoch 2/5
300/300 [==============================] - 7s 25ms/step - loss: 0.0381 - sparse_categorical_accuracy: 0.9878 - val_loss: 0.0379 - val_sparse_categorical_accuracy: 0.9885
Epoch 3/5
300/300 [==============================] - 8s 26ms/step - loss: 0.0213 - sparse_categorical_accuracy: 0.9933 - val_loss: 0.0366 - val_sparse_categorical_accuracy: 0.9892
Epoch 4/5
300/300 [==============================] - 8s 25ms/step - loss: 0.0130 - sparse_categorical_accuracy: 0.9961 - val_loss: 0.0327 - val_sparse_categorical_accuracy: 0.9902
Epoch 5/5
300/300 [==============================] - 8s 25ms/step - loss: 0.0083 - sparse_categorical_accuracy: 0.9975 - val_loss: 0.0465 - val_sparse_categorical_accuracy: 0.9885
<tensorflow.python.keras.callbacks.History at 0x7fc0906c90b8>

Para reducir la sobrecarga de Python y maximizar el rendimiento de su TPU, pruebe el argumento experimental experimental_steps_per_execution en Model.compile . Aquí 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.
                experimental_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)
WARNING:tensorflow:The argument `steps_per_execution` is no longer experimental. Pass `steps_per_execution` instead of `experimental_steps_per_execution`.
WARNING:tensorflow:The argument `steps_per_execution` is no longer experimental. Pass `steps_per_execution` instead of `experimental_steps_per_execution`.
Epoch 1/5
300/300 [==============================] - 14s 46ms/step - loss: 0.2399 - sparse_categorical_accuracy: 0.9279 - val_loss: 0.0510 - val_sparse_categorical_accuracy: 0.9838
Epoch 2/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0406 - sparse_categorical_accuracy: 0.9876 - val_loss: 0.0409 - val_sparse_categorical_accuracy: 0.9882
Epoch 3/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0203 - sparse_categorical_accuracy: 0.9936 - val_loss: 0.0394 - val_sparse_categorical_accuracy: 0.9879
Epoch 4/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0132 - sparse_categorical_accuracy: 0.9959 - val_loss: 0.0410 - val_sparse_categorical_accuracy: 0.9879
Epoch 5/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0140 - sparse_categorical_accuracy: 0.9953 - val_loss: 0.0508 - val_sparse_categorical_accuracy: 0.9863
<tensorflow.python.keras.callbacks.History at 0x7fbc743646d8>

Entrene un modelo usando un ciclo de entrenamiento personalizado.

También puede crear y entrenar sus modelos usando las API tf.function y tf.distribute directamente. strategy.experimental_distribute_datasets_from_function API se utiliza para distribuir el conjunto de datos dada una función de conjunto de datos. Tenga en cuenta que el tamaño de lote que se pasa al conjunto de datos será por tamaño de lote de réplica en lugar del tamaño de lote global en este caso. 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 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-2f075cabff81>: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-2f075cabff81>: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.1237, accuracy: 96.17%
Epoch: 1/5
Current step: 600, training loss: 0.0333, accuracy: 98.98%
Epoch: 2/5
Current step: 900, training loss: 0.0185, accuracy: 99.37%
Epoch: 3/5
Current step: 1200, training loss: 0.0135, accuracy: 99.54%
Epoch: 4/5
Current step: 1500, training loss: 0.0092, accuracy: 99.7%

Mejora del rendimiento mediante varios pasos dentro de tf.function

El rendimiento se puede mejorar ejecutando varios pasos dentro de una función tf.function . Esto se consigue envolviendo la strategy.run llamada con un tf.range dentro tf.function , Autograph lo convertirá en un tf.while_loop en el trabajador de TPU.

Aunque con un mejor rendimiento, existen compensaciones en comparación con 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.009, accuracy: 99.71%

Próximos pasos