API de TPUEstimator en Cloud TPU

En este documento, se explica el uso de la API de TPUEstimator con Cloud TPU. TPUEstimator simplifica la ejecución de modelos en Cloud TPU mediante la administración de numerosos detalles específicos de hardware de bajo nivel.

Los modelos escritos con TPUEstimator funcionan en CPU, GPU, dispositivos de TPU única y pods de TPU, 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, etcétera.

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

En el siguiente código, se muestra cómo entrenar 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 utiliza 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:

En TPUEstimator, la instancia principal 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 de Cloud TPU. Una configuración típica establece la instancia principal de la sesión de TensorFlow como el primer trabajador.
La canalización de entrada se ubica en hosts remotos (en lugar de locales) para garantizar que los ejemplos de entrenamiento se transmitan a Cloud TPU lo más rápido posible. Se requiere un conjunto de datos (tf.data).
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

Recomendamos que primero transfieras un modelo pequeño y pruebes su comportamiento. Esto te permite consolidar tu familiaridad con los conceptos básicos de TPUEstimator. Cuando se ejecute el modelo, agrega más funcionalidad de manera gradual.

Consulta los instructivos para ver un conjunto de modelos de muestra y las instrucciones a fin de 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:

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 canalización de entrada deben ejecutarse en el trabajador para obtener un buen rendimiento. tf.data ejecuta las operaciones en el trabajador.
A fin de amortizar el costo de lanzamiento de TPU, el paso de entrenamiento de modelos se une a un tf.while_loop, en el que un Session.run ejecuta muchas iteraciones para un solo bucle de entrenamiento. Actualmente, solo se puede unir tf.data en un 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.