TPUEstimator API in Cloud TPU

In diesem Dokument wird die Verwendung der Twilio API mit Cloud TPU beschrieben. Twilio vereinfacht das Ausführen von Modellen in einer Cloud TPU durch die Verarbeitung vieler Low-Level-Hardware-spezifischer Details.

Mit Twilio geschriebene Modelle funktionieren in der Regel ohne Codeänderungen auf CPUs, GPUs, einzelnen TPU-Geräten und TPU-Pods. TPUEstimator erleichtert auch das Erreichen maximaler Leistung, da einige Optimierungen automatisch durchgeführt werden.

Informationen zur allgemeinen Funktionsweise von ML-Arbeitslasten auf TPU-Hardware finden Sie in der Dokumentation zur Systemarchitektur.

Standardmäßige TensorFlow Estimator API

Auf übergeordneter Ebene können Sie mit der standardmäßigen TensorFlow Estimator API folgende Schritte ausführen:

Estimator.train() – Ein Modell für eine bestimmte Eingabe in einer festgelegten Anzahl von Schritten trainieren
Estimator.evaluate() – Das Modell an einem Test-Dataset auswerten
Estimator.predict() – Mit dem trainierten Modell eine Inferenz ausführen
Estimator.export_savedmodel() – Das Modell für die Bereitstellung exportieren

Darüber hinaus enthält Estimator ein Standardverhalten für Trainingsjobs, z. B. das Speichern und Wiederherstellen von Prüfpunkten, das Erstellen von Zusammenfassungen für TensorBoard und so weiter.

Estimator erfordert das Schreiben der Funktionen model_fn und input_fn, die den Modell- und Eingabeschnitten des TensorFlow-Diagramms entsprechen.

TPUEstimator-Programmiermodell

TPUEstimator fasst die Berechnung zusammen (model_fn) und verteilt sie auf alle verfügbaren Cloud TPU-Kerne. Die Lernrate muss auf die Batchgröße abgestimmt werden.

Die Funktion input_fn modelliert die Eingabepipeline, die auf der Remote-Host-CPU ausgeführt wird. Mit tf.data können Sie die Eingabevorgänge programmieren, wie in der Programmieranleitung beschrieben. Jeder Aufruf behandelt Eingaben des allgemeinen Batches auf einem Gerät. Die Shard-Batchgröße wird aus params['batch_size'] abgerufen. Profi-Tipp: Lassen Sie ein Dataset anstelle von Tensoren zurückgeben, um eine optimale Leistung zu erzielen.
Die Funktion model_fn modelliert die Berechnung, die repliziert und an die TPUs verteilt wird. Die Berechnung sollte nur Vorgänge enthalten, die von Cloud TPU unterstützt werden. Unter TensorFlow-Vorgänge finden Sie die Liste der verfügbaren Vorgänge.

Trainingsbeispiel mit TPUEstimator

Der folgende Code zeigt das Training von MNIST mit 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)

Der folgende Abschnitt behandelt die neuen Konzepte, die im vorherigen Beispiel eingeführt wurden, damit Sie Cloud TPU effektiv nutzen können.

TPUEstimator-Konzepte

TPUEstimator verwendet für die Ausführung von TensorFlow-Programmen einen "In-Graph"-Replikationsansatz. Die "In-Graph"-Replikation (Einzelsitzung) unterscheidet sich von der "Between-Graph"-Replikation (Multi-Sitzung), die normalerweise in verteilten TensorFlow-Konfigurationen verwendet wird. Dies sind die wichtigsten Unterschiede:

Bei TPUEstimator ist der TensorFlow-Sitzungsinitiator nicht lokal. Ihr Python-Programm erstellt eine einzelne Grafik, die auf allen Kernen in der Cloud TPU repliziert wird. Bei einer typischen Konfiguration wird der TensorFlow-Sitzungsinitiator als erster Worker festgelegt.
Die Eingabepipeline wird auf Remote-Hosts (statt lokal) platziert, um sicherzustellen, dass Trainingsbeispiele so schnell wie möglich in die Cloud TPU eingespeist werden. Ein Dataset (tf.data) ist erforderlich.
Cloud TPU-Worker arbeiten synchron, wobei jeder Worker den gleichen Schritt zur gleichen Zeit ausführt.

Von TensorFlow Estimator in TPUEstimator konvertieren

Wir empfehlen, zuerst ein kleines Modell zu portieren und sein Verhalten zu testen. Auf diese Weise können Sie sich ausgiebig mit den grundlegenden Konzepten von TPUEstimator vertraut machen. Wenn Ihr Modell ausgeführt wird, fügen Sie nach und nach weitere Funktionen hinzu.

In den Anleitungen finden Sie eine Reihe von Beispielmodellen und eine Anleitung für die Ausführung mit Cloud TPU. Weitere Modelle sind auf GitHub verfügbar.

Ändern Sie Folgendes, um den Code der Klasse tf.estimator.Estimator für die Verwendung von tf.contrib.tpu.TPUEstimator zu konvertieren:

Ändern Sie tf.estimator.RunConfig zu tf.contrib.tpu.RunConfig.
Legen Sie TPUConfig (Teil von tf.contrib.tpu.RunConfig) fest, um iterations_per_loop anzugeben. iterations_per_loop ist die Anzahl der Iterationen, die auf der Cloud TPU ausgeführt werden sollen session.run (pro Trainingsschleife).

Die Cloud TPU führt eine bestimmte Anzahl von Iterationen der Trainingsschleife aus, bevor sie zum Host zurückkehrt. Prüfpunkte oder Zusammenfassungen werden erst gespeichert, wenn alle Cloud TPU-Iterationen ausgeführt wurden.

Verwenden Sie in model_fn tf.contrib.tpu.CrossShardOptimizer, um die Optimierung zusammenzufassen. Beispiel:

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

Ändern Sie tf.estimator.Estimator zu tf.contrib.tpu.TPUEstimator.

Die standardmäßige RunConfig speichert Zusammenfassungen für TensorBoard alle 100 Schritte und schreibt Prüfpunkte alle 10 Minuten.

FAQ

Warum ist `tf.data` für die Eingabepipeline erforderlich?

Aus zwei Gründen:

Ihr Anwendungscode läuft auf dem Client, die TPU-Berechnung wird jedoch auf dem worker ausgeführt. Vorgänge der Eingabe-Pipeline müssen auf dem Worker ausgeführt werden, um eine gute Leistung zu erzielen. tf.data führt die Vorgänge auf dem Worker aus.
Zum Amortisieren der TPU-Startkosten wird der Modelltrainingsschritt in eine tf.while_loop eingebunden, wobei ein Session.run viele Iterationen für eine einzelne Trainingsschleife ausführt. Derzeit kann nur tf.data in eine tf.while_loop verpackt werden.

Wie kann ich für Leistung des Modelltrainings ein Profil erstellen?

Sie können ein Profil der Modelltrainingsleistung mit dem für TensorBoard bereitgestellten Profiler erstellen.