API TPUEstimator sur Cloud TPU

Ce document couvre l'utilisation de l'API TPUEstimator avec Cloud TPU. TPUEstimator simplifie l'exécution de modèles sur un Cloud TPU en gérant de nombreux détails spécifiques, de bas niveau et spécifiques au matériel.

Les modèles écrits à l'aide de TPUEstimator fonctionnent sur les processeurs, les GPU, les appareils TPU uniques et les pods TPU, généralement sans changement de code. En outre, TPUEstimator effectue automatiquement certaines optimisations à votre place pour vous permettre d'atteindre plus facilement des performances maximales.

Pour en savoir plus sur le fonctionnement des charges de travail de machine learning sur le matériel TPU en général, consultez la documentation sur l'architecture du système.

API Estimator TensorFlow standard

Dans le cadre d'une utilisation générale, l'API Estimator TensorFlow standard fournit les éléments suivants :

Estimator.train() : entraîne un modèle sur une entrée donnée pour un nombre déterminé de pas.
Estimator.evaluate() : évalue le modèle sur un ensemble de test.
Estimator.predict() : exécute l'inférence à l'aide du modèle entraîné.
Estimator.export_savedmodel() : exporte votre modèle pour le diffuser.

En outre, Estimator inclut le comportement par défaut commun aux tâches d'entraînement, telles que l'enregistrement et la restauration de points de contrôle, la création de résumés pour TensorBoard, etc.

Estimator nécessite la spécification de fonctions model_fn et input_fn qui correspondent aux portions de modèle et de saisie de votre graphe TensorFlow.

Modèle de programmation TPUEstimator

Le TPUEstimator encapsule le calcul (le model_fn) et le distribue à tous les cœurs Cloud TPU disponibles. Le taux d'apprentissage doit être ajusté en fonction de la taille de lot.

La fonction input_fn modélise le pipeline d'entrée qui s'exécute sur le processeur hôte distant. Utilisez tf.data pour programmer les opérations de saisie comme décrit dans le guide du programmeur. Chaque appel traite l'entrée du lot global sur un appareil. La taille du lot de partitionnement est récupérée à partir de params['batch_size']. Conseil : renvoyez un ensemble de données plutôt que des Tensors pour optimiser les performances.
La fonction model_fn modélise le calcul à répliquer et à distribuer aux TPU. Le calcul ne doit contenir que des opérations compatibles avec Cloud TPU. Les opérations TensorFlow incluent la liste des opérations disponibles.

Exemple d'apprentissage à l'aide de TPUEstimator

Le code suivant illustre l'entraînement MNIST avec 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)

La section suivante aborde les nouveaux concepts introduits dans l'exemple précédent pour vous aider à utiliser efficacement Cloud TPU.

Concepts relatifs à TPUEstimator

TPUEstimator adopte une approche de réplication intra-graphe lors de l'exécution de programmes TensorFlow. La réplication intra-graphe (session unique) diffère de l'apprentissage par réplication entre plusieurs graphes (sessions multiples) généralement utilisé sur la plate-forme TensorFlow distribuée. Voici les principales différences entre ces deux approches :

Dans TPUEstimator, l'initiateur de session TensorFlow n'est pas local. Votre programme Python crée un graphe unique qui est répliqué sur tous les cœurs du Cloud TPU. L'initiateur de session TensorFlow est généralement défini en tant que premier nœud de calcul.
Le pipeline d'entrée est placé sur des hôtes distants (et non sur des instances locales) afin d'envoyer les exemples d'entraînement à Cloud TPU aussi rapidement que possible. Vous devez indiquer un ensemble de données tf.data.
Les nœuds de calcul Cloud TPU fonctionnent de manière synchrone (ils effectuent tous la même étape au même moment).

Convertir un modèle Estimator TensorFlow au format TPUEstimator

Nous vous recommandons de commencer par transférer un petit modèle et de tester son comportement. Cette démarche vous aide à mieux comprendre les concepts de base de TPUEstimator. Lorsque votre modèle s'exécute, ajoutez progressivement des fonctionnalités supplémentaires.

Consultez les tutoriels pour obtenir un ensemble d'exemples de modèles et d'instructions pour les exécuter avec Cloud TPU. Des modèles supplémentaires sont disponibles sur GitHub.

Pour convertir votre code de la classe tf.estimator.Estimator et utiliser la classe tf.contrib.tpu.TPUEstimator, modifiez les éléments suivants :

Passer de tf.estimator.RunConfig à tf.contrib.tpu.RunConfig.
Définissez TPUConfig (une partie de tf.contrib.tpu.RunConfig) pour spécifier iterations_per_loop. iterations_per_loop correspond au nombre d'itérations à exécuter sur Cloud TPU pour un appel de session.run (par boucle d'apprentissage).

Cloud TPU exécute un certain nombre d'itérations de la boucle d'apprentissage avant de revenir à l'hôte. Aucun point de contrôle ni résumé n'est enregistré avant que toutes les itérations Cloud TPU n'aient été exécutées.

Dans model_fn, utilisez tf.contrib.tpu.CrossShardOptimizer pour recouvrir votre optimiseur. Exemple :

 optimizer = tf.contrib.tpu.CrossShardOptimizer(
      tf.train.GradientDescentOptimizer(learning_rate=learning_rate))

Passer de tf.estimator.Estimator à tf.contrib.tpu.TPUEstimator.

La valeur par défaut RunConfig enregistre les résumés pour TensorBoard tous les 100 pas et écrit des points de contrôle toutes les 10 minutes.

Questions fréquentes

Pourquoi le pipeline d'entrée nécessite-t-il `tf.data` ?

Il existe deux raisons :

Votre code d'application s'exécute sur le client pendant que le calcul TPU est exécuté sur le worker. Les opérations du pipeline d'entrée doivent être exécutées sur le nœud de calcul pour obtenir de bonnes performances. tf.data exécute les opérations sur le nœud de calcul.
Afin d'amortir le coût de lancement du TPU, l'étape d'entraînement du modèle est encapsulée dans une tf.while_loop, où un objet Session.run exécute de nombreuses itérations pour une seule boucle d'entraînement. Actuellement, seul tf.data peut être encapsulé dans un tf.while_loop.

Comment puis-je profiler les performances de l'entraînement de modèle ?

Vous pouvez profiler les performances de l'entraînement du modèle à l'aide du profileur fourni pour TensorBoard.