Usa la API de Estimator de TPU en Cloud TPU

En este documento, se describe el uso de la API de TPUEstimator con Cloud TPU. TPUEstimator simplifica la ejecución de modelos en Cloud TPU; para ello, maneja varios detalles específicos de hardware de nivel bajo.

Los modelos escritos con TPUEstimator funcionan en CPU, GPU, dispositivos de TPU única y pods de TPU completos, por lo general, sin cambios en el código. TPUEstimator también facilita la obtención del máximo rendimiento; para ello, realiza de forma automática algunas optimizaciones en tu nombre.

Para saber cómo funcionan las cargas de trabajo de aprendizaje automático en el hardware de TPU en general, lee la documentación de Arquitectura del sistema.

API de Estimator de TensorFlow estándar

En los niveles altos, la API de Estimator de TensorFlow estándar establece lo siguiente:

  • Estimator.train(): entrena un modelo en una entrada determinada durante una cantidad fija de pasos.
  • Estimator.evaluate(): evalúa el modelo en un conjunto de prueba.
  • Estimator.predict(): ejecuta inferencia con el modelo entrenado.
  • Estimator.export_savedmodel(): exporta tu modelo para la entrega.

Además, Estimator incluye el comportamiento predeterminado común a los trabajos de entrenamiento, como guardar y restablecer puntos de control, crear resúmenes para TensorBoard, entre otras opciones.

Estimator requiere que escribas model_fn y input_fn que correspondan al modelo y las porciones de entrada de tu grafo de TensorFlow.

Modelo de programación de TPUEstimator

TPUEstimator une el cálculo (model_fn) y lo distribuye a todos los núcleos de Cloud TPU disponibles. La tasa de aprendizaje debe ajustarse según el tamaño del lote.

  • La función input_fn modela la canalización de entrada que se ejecuta en la CPU host remota. Usa tf.data para programar las operaciones de entrada como se describe en la guía del programador. Cada invocación maneja la entrada del lote global en un dispositivo. El tamaño del lote del fragmento se recupera de params['batch_size']. Sugerencia profesional: muestra un conjunto de datos en lugar de tensores para obtener un rendimiento óptimo.

  • La función model_fn modela el cálculo que se replica y distribuye a las TPU. El cálculo debe contener solo operaciones que Cloud TPU admite. Las operaciones de TensorFlow incluyen la lista de operaciones disponibles.

Un ejemplo de entrenamiento con TPUEstimator

El siguiente código muestra cómo entrenar un modelo MNIST con TPUEstimator:

def model_fn(features, labels, mode, params):
  """A simple CNN."""
  del params  # unused

  input_layer = tf.reshape(features, [-1, 28, 28, 1])
  conv1 = tf.layers.conv2d(
      inputs=input_layer, filters=32, kernel_size=[5, 5], padding="same",
      activation=tf.nn.relu)
  pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)
  conv2 = tf.layers.conv2d(
      inputs=pool1, filters=64, kernel_size=[5, 5],
      padding="same", activation=tf.nn.relu)
  pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)
  pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
  dense = tf.layers.dense(inputs=pool2_flat, units=128, activation=tf.nn.relu)
  dropout = tf.layers.dropout(
      inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
  logits = tf.layers.dense(inputs=dropout, units=10)
  onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10)

  loss = tf.losses.softmax_cross_entropy(
      onehot_labels=onehot_labels, logits=logits)

  learning_rate = tf.train.exponential_decay(
      FLAGS.learning_rate, tf.train.get_global_step(), 100000, 0.96)

  optimizer = tpu_optimizer.CrossShardOptimizer(
      tf.train.GradientDescentOptimizer(learning_rate=learning_rate))

  train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
  return tpu_estimator.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op)

def make_input_fn(filename):
  """Returns an `input_fn` for train and eval."""

  def input_fn(params):
    """An input_fn to parse 28x28 images from filename using tf.data."""
    batch_size = params["batch_size"]

    def parser(serialized_example):
      """Parses a single tf.Example into image and label tensors."""
      features = tf.parse_single_example(
          serialized_example,
          features={
              "image_raw": tf.FixedLenFeature([], tf.string),
              "label": tf.FixedLenFeature([], tf.int64),
          })
      image = tf.decode_raw(features["image_raw"], tf.uint8)
      image.set_shape([28 * 28])
      # Normalize the values of the image from the range [0, 255] to [-0.5, 0.5]
      image = tf.cast(image, tf.float32) * (1. / 255) - 0.5
      label = tf.cast(features["label"], tf.int32)
      return image, label

    dataset = tf.contrib.data.TFRecordDataset(
        filename, buffer_size=FLAGS.dataset_reader_buffer_size)
    dataset = dataset.repeat()
    dataset = dataset.apply(
      tf.contrib.data.map_and_batch(
         parser, batch_size=batch_size,
         num_parallel_batches=8,
         drop_remainder=True))
    return dataset

  return input_fn

def main(unused_argv):

  tf.logging.set_verbosity(tf.logging.INFO)

  run_config = tpu_config.RunConfig(
      master=FLAGS.master,
      model_dir=FLAGS.model_dir,
      session_config=tf.ConfigProto(
          allow_soft_placement=True, log_device_placement=True),
      tpu_config=tpu_config.TPUConfig(FLAGS.iterations))

  estimator = tpu_estimator.TPUEstimator(
      model_fn=model_fn,
      use_tpu=FLAGS.use_tpu,
      train_batch_size=FLAGS.batch_size,
      eval_batch_size=FLAGS.batch_size,
      config=run_config)

  estimator.train(input_fn=make_input_fn(FLAGS.train_file),
                  max_steps=FLAGS.train_steps)

En la siguiente sección, se abordan los conceptos nuevos ingresados en el ejemplo anterior para ayudarte a usar Cloud TPU de manera efectiva.

Conceptos de TPUEstimator

TPUEstimator usa un enfoque de replicación en el grafo para ejecutar programas de TensorFlow. La replicación en el grafo (sesión única) difiere del entrenamiento de replicación entre grafos (sesiones múltiples) que se suele utilizar en TensorFlow distribuido. A continuación, se mencionan las diferencias principales:

  1. En TPUEstimator, el iniciador de la sesión de TensorFlow no es local. Tu programa de Python crea un grafo único que se replica en todos los núcleos en Cloud TPU. Una configuración típica establece el iniciador de la sesión de TensorFlow para que sea el primer trabajador.

  2. La canalización de entrada se coloca en hosts remotos (en lugar de locales) para garantizar que los ejemplos de entrenamiento puedan enviarse a Cloud TPU tan rápido como sea posible. Se requiere un conjunto de datos (tf.data).

  3. Los trabajadores de Cloud TPU operan de forma síncrona y cada uno realiza el mismo paso al mismo tiempo.

Realizar la conversión de TensorFlow Estimator a TPUEstimator

Te recomendamos que primero cargues un modelo pequeño y simple y pruebes el comportamiento de extremo a extremo. Esto te permite consolidar tu familiaridad con los conceptos básicos de TPUEstimator. Cuando se ejecute el modelo simple, agrega gradualmente más funcionalidad.

Consulta los instructivos a fin de obtener un conjunto de modelos de muestra y, además, instrucciones para ejecutarlos con Cloud TPU. Hay modelos adicionales disponibles en GitHub.

A fin de convertir tu código de la clase tf.estimator.Estimator para usar tf.contrib.tpu.TPUEstimator, haz esto:

  • Cambia tf.estimator.RunConfig a tf.contrib.tpu.RunConfig.
  • Establece TPUConfig (parte de tf.contrib.tpu.RunConfig) para especificar iterations_per_loop. iterations_per_loop es la cantidad de iteraciones que se ejecutarán en Cloud TPU para una llamada session.run (por ciclo de entrenamiento).

Cloud TPU ejecuta un número específico de iteraciones del ciclo de entrenamiento antes de volver al host. No se guardan puntos de control ni resúmenes hasta que se ejecuten todas las iteraciones de Cloud TPU.

  • En model_fn, usa tf.contrib.tpu.CrossShardOptimizer para ajustar tu optimizador. Por ejemplo:

     optimizer = tf.contrib.tpu.CrossShardOptimizer(
          tf.train.GradientDescentOptimizer(learning_rate=learning_rate))
    
  • Cambia tf.estimator.Estimator a tf.contrib.tpu.TPUEstimator.

La RunConfig predeterminada guarda resúmenes para TensorBoard cada 100 pasos y escribe puntos de control cada 10 minutos.

Preguntas frecuentes

¿Por qué se requiere tf.data para la canalización de entrada?

Existen dos motivos:

  1. El código de la aplicación se ejecuta en el cliente, mientras que el cálculo de TPU se ejecuta en el worker. Las operaciones de la canalización de entrada deben colocarse en el trabajador remoto para obtener un buen rendimiento. Solo tf.data (conjunto de datos) es compatible con esto.

  2. Para amortizar el costo de lanzamiento de TPU, el paso de entrenamiento de modelos se une a un tf.while_loop, de manera que una Session.run ejecuta muchas iteraciones para un solo ciclo de entrenamiento. Actualmente, solo un tf.data puede unirse a tf.while_loop.

¿Cómo creo el perfil de rendimiento del entrenamiento de modelos?

Puede crear el perfil de rendimiento del entrenamiento de modelos con el generador de perfiles provisto para TensorBoard.