¡Confirme su asistencia a su evento local de TensorFlow Everywhere hoy!
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 actualmente está disponible 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.

Preparar

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='grpc://' + os.environ['COLAB_TPU_ADDR'])
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.58:8470

INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.58: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. Esta guía demuestra 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 Cloud TPU, ya que es imposible usar las Cloud TPU 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 (usando 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. Vea 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í debería 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 usara una TPUStrategy MirroredStrategy lugar de una 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.3140 - sparse_categorical_accuracy: 0.8968 - val_loss: 0.0428 - val_sparse_categorical_accuracy: 0.9852
Epoch 2/5
300/300 [==============================] - 8s 26ms/step - loss: 0.0399 - sparse_categorical_accuracy: 0.9876 - val_loss: 0.0414 - val_sparse_categorical_accuracy: 0.9865
Epoch 3/5
300/300 [==============================] - 7s 25ms/step - loss: 0.0204 - sparse_categorical_accuracy: 0.9934 - val_loss: 0.0488 - val_sparse_categorical_accuracy: 0.9864
Epoch 4/5
300/300 [==============================] - 8s 25ms/step - loss: 0.0141 - sparse_categorical_accuracy: 0.9956 - val_loss: 0.0674 - val_sparse_categorical_accuracy: 0.9819
Epoch 5/5
300/300 [==============================] - 8s 25ms/step - loss: 0.0104 - sparse_categorical_accuracy: 0.9965 - val_loss: 0.0449 - val_sparse_categorical_accuracy: 0.9885

<tensorflow.python.keras.callbacks.History at 0x7fede8002940>

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 48ms/step - loss: 0.2434 - sparse_categorical_accuracy: 0.9269 - val_loss: 0.0461 - val_sparse_categorical_accuracy: 0.9861
Epoch 2/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0378 - sparse_categorical_accuracy: 0.9884 - val_loss: 0.0394 - val_sparse_categorical_accuracy: 0.9877
Epoch 3/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0207 - sparse_categorical_accuracy: 0.9935 - val_loss: 0.0484 - val_sparse_categorical_accuracy: 0.9861
Epoch 4/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0145 - sparse_categorical_accuracy: 0.9950 - val_loss: 0.0406 - val_sparse_categorical_accuracy: 0.9878
Epoch 5/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0105 - sparse_categorical_accuracy: 0.9967 - val_loss: 0.0528 - val_sparse_categorical_accuracy: 0.9857

<tensorflow.python.keras.callbacks.History at 0x7fe9242485f8>

Entrene un modelo utilizando 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.1293, accuracy: 95.98%
Epoch: 1/5
Current step: 600, training loss: 0.034, accuracy: 98.94%
Epoch: 2/5
Current step: 900, training loss: 0.0182, accuracy: 99.39%
Epoch: 3/5
Current step: 1200, training loss: 0.0143, accuracy: 99.53%
Epoch: 4/5
Current step: 1500, training loss: 0.009, accuracy: 99.69%

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 logra tf.range llamada a strategy.run con un tf.range dentro de 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.0066, accuracy: 99.78%

Próximos pasos