Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
tensorflow
GitHub Repository: tensorflow/docs-l10n
 Path: blob/master/site/es-419/guide/tpu.ipynb
25115 views
Kernel: Python 3
#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

En esta guía se muestra cómo realizar un entrenamiento básico en las Unidades de Procesamiento de Tensores (TPUs) y TPU Pods, una colección de dispositivos TPU conectados por interfaces de red de alta velocidad especializadas, con tf.keras y bucles de entrenamiento personalizados.

Las TPU son circuitos integrados de aplicación específica (ASIC) desarrollados a medida por Google que se utilizan para acelerar las cargas de trabajo de aprendizaje automático. Están disponibles en Google Colab, TPU Research Cloud y Cloud TPU.

Preparación

Antes de ejecutar este bloc de notas de Colab, asegúrese de que su acelerador de hardware es un TPU verificando la configuración del bloc de notas: Runtime > Change runtime type > Hardware accelerator > TPU.

Importe algunas bibliotecas necesarias, incluyendo los conjuntos de datos de TensorFlow:

import tensorflow as tf

import os
import tensorflow_datasets as tfds

Inicialización de la TPU

Las TPUs normalmente son TPUs en la nube que trabajan, las cuales son diferentes del proceso local que ejecuta el programa Python del usuario. Por lo tanto, es necesario hacer algún trabajo de inicialización para conectarse al conjunto remoto e inicializar las TPUs. Tenga en cuenta que el argumento tpu de tf.distribute.cluster_resolver.TPUClusterResolver es una dirección especial sólo para Colab. Si está ejecutando su código en Google Compute Engine (GCE), debe pasar en su lugar el nombre de su TPU en la nube.

Nota: El código de inicialización del TPU tiene que estar al principio de su programa.

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'))

Colocación manual en el dispositivo

Después de inicializar la TPU, puede utilizar la colocación manual de dispositivos para situar el cálculo en un único dispositivo 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)

Estrategias de distribución

Normalmente, su modelo se ejecuta en múltiples TPUs de forma paralela a los datos. Para distribuir su modelo en múltiples TPUs (así como múltiples GPUs o múltiples máquinas), TensorFlow ofrece la API tf.distribute.Strategy. Puede reemplazar su estrategia de distribución y el modelo se ejecutará en cualquier dispositivo (TPU) dado. Obtenga más información en la guía Entrenamiento distribuido con TensorFlow.

El uso de la opción tf.distribute.TPUStrategy implementa el entrenamiento distribuido sincronizado. Las TPUs proporcionan su propia implementación de operaciones eficientes de reducción total y otras operaciones colectivas entre varios núcleos de TPU, que se utilizan en TPUStrategy.

Para demostrarlo, cree un objeto tf.distribute.TPUStrategy:

strategy = tf.distribute.TPUStrategy(resolver)

Para replicar un cálculo de modo que pueda ejecutarse en todos los núcleos de la TPU, puede transferirlo a la API Strategy.run. A continuación se muestra un ejemplo en el que todos los núcleos 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)

Clasificación en TPUs

Después de haber cubierto los conceptos básicos, considere un ejemplo más concreto. Esta sección muestra cómo utilizar la estrategia de distribución -tf.distribute.TPUStrategy- para entrenar un modelo Keras en una TPU en la nube.

Definir un modelo Keras

Comenzamos con la definición de un modelo Sequential de Keras para la clasificación de imágenes en el conjunto de datos MNIST. No es diferente de lo que utilizaría si estuviera entrenando en CPUs o GPUs. Tenga en cuenta que la creación del modelo Keras tiene que estar dentro del Strategy.scope, por lo que las variables se pueden crear en cada dispositivo TPU. Otras partes del código no necesitan estar dentro del ámbito Strategy.

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

Este modelo pone términos de regularización L2 en los pesos de cada capa, de modo que el bucle de entrenamiento personalizado de abajo puede mostrar cómo los recoge de Model.losses.

Carga del conjunto de datos.

El uso eficiente de la API tf.data.Dataset es fundamental cuando se utiliza una TPU en la nube. Puede obtener más información sobre el rendimiento de los conjuntos de datos en la Guía de rendimiento respecto a las canalizaciones de entrada.

Si está utilizando Nodos UTP, necesita almacenar todos los archivos de datos leídos por el Dataset de TensorFlow en cubos de Google Cloud Storage (GCS). Si utiliza TPU VMs, puede almacenar los datos donde desee. Para obtener más información sobre los nodos TPU y las máquinas virtuales TPU, consulte la documentación Arquitectura del sistema TPU.

Para la mayoría de los casos de uso, se recomienda convertir los datos en el formato TFRecord y utilizar un tf.data.TFRecordDataset para leerlos. Consulte el tutorial TFRecord y tf.Example para obtener más detalles sobre cómo hacerlo. No es un requisito obligatorio y puede utilizar otros lectores de conjuntos de datos, como tf.data.FixedLengthRecordDataset o tf.data.TextLineDataset.

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

Independientemente del formato de datos utilizado, se recomienda encarecidamente utilizar archivos de gran tamaño, del orden de 100 MB. Esto es particularmente importante en esta configuración en red, ya que la sobrecarga de abrir un archivo es significativamente mayor.

Como se muestra en el siguiente código, debe utilizar el módulo de conjuntos de datos de Tensorflow tfds.load para obtener una copia de los datos de entrenamiento y prueba MNIST. Tenga en cuenta que try_gcs es específico para utilizar una copia que esté disponible en un bucket GCS público. Si no se 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 utilizando las API de alto nivel de Keras

Puede entrenar su modelo con las API de Keras Model.fit y Model.compile. . No hay nada específico de TPU en este paso: escriba el código si utiliza múltiples GPU y una MirroredStrategy en vez de 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)

Para reducir la sobrecarga de Python y maximizar el rendimiento de su TPU, pase el argumento steps_per_execution al comando Model.compile de Keras. En este ejemplo, se 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)

Entrene el modelo utilizando un bucle de entrenamiento personalizado

También puede crear y entrenar su modelo utilizando tf.function y tf.distribute directamente. Puede utilizar la API Strategy.experimental_distribute_datasets_from_function para distribuir el tf.data.Dataset dado una función de conjunto de datos. Tenga en cuenta que en el siguiente ejemplo el tamaño del lote que se pasa al Dataset es el tamaño del lote por réplica en vez del tamaño del lote global. Para obtener más información, consulte el tutorial Entrenamiento personalizado con tf.distribute.Strategy.

En primer lugar, cree el modelo, los conjuntos de datos y las tf.function:

# Create the model, optimizer and metrics inside the `tf.distribute.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 `tf.data.Dataset`s
# 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)
      per_example_loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(per_example_loss)
      model_losses = model.losses
      if model_losses:
        loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))

    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),))

Después, ejecute el bucle 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: {}, training 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()

Cómo mejorar el rendimiento con varios pasos dentro de tf.function

Puede mejorar el rendimiento ejecutando varios pasos dentro de una tf.function. Esto se consigue envolviendo la llamada Strategy.run con un tf.range dentro de tf.function, y AutoGraph lo convertirá en un tf.while_loop en el trabajador de la TPU. Puede obtener más información sobre tf.functions en la guía Mejorar el rendimiento con tf.function.

A pesar de mejorar el rendimiento, hay desventajas con este método en comparación con la ejecución de un solo paso dentro de un tf.function. Ejecutar múltiples pasos en un tf.function es menos flexible, no puede ejecutar cosas eagerly o un código 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)
      per_example_loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(per_example_loss)
      model_losses = model.losses
      if model_losses:
        loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))
    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: {}, training accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))

Siguientes pasos

Para obtener más información sobre las TPU en la nube y cómo utilizarlas: