Usar TPU

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

Antes de ejecutar este portátil Colab, asegúrese de que el acelerador de hardware es un TPU comprobando la configuración de portátiles: 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
/home/kbuilder/.local/lib/python3.6/site-packages/requests/__init__.py:104: RequestsDependencyWarning: urllib3 (1.26.6) or chardet (2.3.0)/charset_normalizer (2.0.4) doesn't match a supported version!
  RequestsDependencyWarning)

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 tpu argumento para tf.distribute.cluster_resolver.TPUClusterResolver es una dirección especial 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:Clearing out eager caches
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2:8470
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
All devices:  [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'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:3', 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:5', 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:7', 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. Compruebe la guía de estrategia de distribución para más información.

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

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 por lo que se puede ejecutar en todos los núcleos de TPU, puede pasar que en el strategy.run API. A continuación se muestra un ejemplo que muestra todos los núcleos de la recepción de las mismas entradas (a, b) y la realización de 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 utilizar la distribución estrategia- tf.distribute.TPUStrategy : para entrenar un modelo Keras en una nube de TPU.

Definir un modelo de Keras

Comenzar con una definición de un Sequential modelo Keras para la clasificación de imágenes en el conjunto de datos utilizando MNIST Keras. No es diferente de lo que usaría si estuviera entrenando en CPU o GPU. Tenga en cuenta que las necesidades de creación de modelos Keras 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 tf.data.Dataset API es crítico cuando se utiliza una nube de TPU, ya que es imposible utilizar la nube TPU a menos que pueda alimentarlos datos con la suficiente rapidez. Se puede obtener más información sobre el rendimiento conjunto de datos en la guía de rendimiento de la tubería de entrada .

Para todos, pero los experimentos simples (utilizando tf.data.Dataset.from_tensor_slices u otros datos en el gráfico), es necesario almacenar todos los archivos de datos leídos por el conjunto de datos en Google Cloud Storage (GCS) cubos.

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

Es posible cargar todo los conjuntos de datos pequeños 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 siguiente, debe utilizar los tensorflow_datasets módulo para obtener una copia de la formación MNIST y datos de prueba. Tenga en cuenta que try_gcs se especifica para utilizar una copia que está disponible en un cubo GCS pública. 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

Usted puede entrenar a su modelo con Keras fit y compile API. No hay nada de TPU-específica en este paso se escribe el código como si estuviera usando GPU mutliple y un MirroredStrategy en lugar de la TPUStrategy . Se puede obtener más información en la formación distribuida con Keras tutorial.

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 35ms/step - loss: 0.1352 - sparse_categorical_accuracy: 0.9583 - val_loss: 0.0404 - val_sparse_categorical_accuracy: 0.9863
Epoch 2/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0346 - sparse_categorical_accuracy: 0.9893 - val_loss: 0.0376 - val_sparse_categorical_accuracy: 0.9882
Epoch 3/5
300/300 [==============================] - 7s 22ms/step - loss: 0.0199 - sparse_categorical_accuracy: 0.9934 - val_loss: 0.0432 - val_sparse_categorical_accuracy: 0.9872
Epoch 4/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0132 - sparse_categorical_accuracy: 0.9957 - val_loss: 0.0475 - val_sparse_categorical_accuracy: 0.9870
Epoch 5/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0100 - sparse_categorical_accuracy: 0.9970 - val_loss: 0.0443 - val_sparse_categorical_accuracy: 0.9877
<keras.callbacks.History at 0x7faf8829f278>

Para reducir los gastos generales y Python maximizar el rendimiento de su TPU, pasar en el argumento, steps_per_execution -para 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 [==============================] - 13s 43ms/step - loss: 0.1405 - sparse_categorical_accuracy: 0.9571 - val_loss: 0.0438 - val_sparse_categorical_accuracy: 0.9860
Epoch 2/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0337 - sparse_categorical_accuracy: 0.9890 - val_loss: 0.0503 - val_sparse_categorical_accuracy: 0.9847
Epoch 3/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0170 - sparse_categorical_accuracy: 0.9947 - val_loss: 0.0448 - val_sparse_categorical_accuracy: 0.9884
Epoch 4/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0137 - sparse_categorical_accuracy: 0.9956 - val_loss: 0.0362 - val_sparse_categorical_accuracy: 0.9903
Epoch 5/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0090 - sparse_categorical_accuracy: 0.9971 - val_loss: 0.0458 - val_sparse_categorical_accuracy: 0.9880
<keras.callbacks.History at 0x7fae1054f908>

Entrena el modelo usando un ciclo de entrenamiento personalizado

También puede crear y entrenar a su modelo usando tf.function y tf.distribute directamente API. Se puede utilizar el strategy.experimental_distribute_datasets_from_function API 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, echa un vistazo a la formación de encargo con tf.distribute.Strategy tutorial.

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.1352, accuracy: 95.85%
Epoch: 1/5
Current step: 600, training loss: 0.0355, accuracy: 98.86%
Epoch: 2/5
Current step: 900, training loss: 0.0194, accuracy: 99.37%
Epoch: 3/5
Current step: 1200, training loss: 0.0127, accuracy: 99.6%
Epoch: 4/5
Current step: 1500, training loss: 0.0111, accuracy: 99.65%

Mejorar el rendimiento con múltiples pasos en el interior tf.function

Puede mejorar el rendimiento mediante la ejecución de varios pasos dentro de un tf.function . Esto se consigue envolviendo la strategy.run llamada con un tf.range dentro tf.function y AutoGraph lo convertirá en un tf.while_loop en el trabajador de TPU.

A pesar de la mejora del rendimiento, hay ventajas y desventajas con este método en comparación con la ejecución de un solo paso en el interior tf.function . La ejecución de varios pasos en un tf.function es menos flexible, no se puede ejecutar cosas con entusiasmo o código Python arbitraria 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.0081, accuracy: 99.73%

Próximos pasos