API TPUEstimator su Cloud TPU

Questo documento illustra l'utilizzo dell'API TPUEstimator con Cloud TPU. TPUEstimator semplifica l'esecuzione di modelli su una Cloud TPU gestendo numerosi dettagli specifici di basso livello e hardware.

I modelli scritti utilizzando TPUEstimator funzionano su CPU, GPU, singoli dispositivi TPU e pod TPU, in genere senza modifiche al codice. TPUEstimator rende inoltre più facile il massimo rendimento eseguendo automaticamente alcune ottimizzazioni per tuo conto.

Per informazioni sul funzionamento generale dei carichi di lavoro di machine learning su hardware TPU, leggi la documentazione sull'architettura di sistema.

API TensorFlow Estimator standard

A livello generale, l'API TensorFlow Estimator fornisce:

Estimator.train(): addestra un modello su un determinato input per un numero fisso di passaggi.
Estimator.evaluate(): valuta il modello in un set di test.
Estimator.predict(): esegui l'inferenza utilizzando il modello addestrato.
Estimator.export_savedmodel(): esporta il modello per la pubblicazione.

Inoltre, Estimator include il comportamento predefinito comune per i job di addestramento, come il salvataggio e il ripristino dei punti di controllo, la creazione di riepiloghi per TensorBoard e così via.

Estimator richiede la scrittura di un valore model_fn e di un input_fn che corrisponda al modello e alle parti di input del grafico TensorFlow.

Modello di programmazione TPUEstimator

TPUEstimator aggrega il calcolo (model_fn) e lo distribuisce a tutti i core Cloud TPU disponibili. La tariffa di apprendimento deve essere regolata in base alle dimensioni del batch.

La funzione input_fn modella la pipeline di input in esecuzione sulla CPU dell'host remoto. Utilizza tf.data per programmare le operazioni di input come descritto nella guida al programma. Ogni chiamata gestisce l'input del batch globale su un dispositivo. Le dimensioni del batch di shard vengono recuperate da params['batch_size']. Suggerimento avanzato: restituisci un set di dati anziché i tensori per prestazioni ottimali.
La funzione model_fn modella il calcolo replicato e distribuito alle TPU. Il calcolo deve contenere solo operazioni supportate da Cloud TPU. Operazioni TensorFlow include l'elenco delle operazioni disponibili.

Esempio di addestramento con TPUEstimator

Il seguente codice dimostra un addestramento 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)

La sezione seguente illustra i nuovi concetti introdotti nell'esempio precedente per aiutarti a utilizzare Cloud TPU in modo efficace.

Concetti relativi a TPUEstimator

TPUEstimator utilizza un approccio di replica nel grafico per eseguire i programmi TensorFlow. La replica della replica all'interno del grafico (a sessione singola) varia rispetto all'addestramento della replica tra il grafico (a più sessioni) in genere utilizzata in TensorFlow distribuito. Le differenze principali includono:

In TPUEstimator, l'iniziatore della sessione TensorFlow non è locale. Il tuo programma Python crea un singolo grafico che viene replicato in tutti i core in Cloud TPU. Una configurazione tipica imposta l'iniziatore di sessione TensorFlow come primo worker.
La pipeline di input viene ospitata su host remoti (anziché locali) per garantire che gli esempi di addestramento vengano inviati alla Cloud TPU il più velocemente possibile. È richiesto un set di dati (tf.data).
I worker di Cloud TPU funzionano in modo sincrono, con ogni worker che esegue lo stesso passaggio contemporaneamente.

Conversione da TensorFlow Estimator a TPUEstimator

Ti consigliamo di trasferire prima un modello di piccole dimensioni e di verificarne il comportamento. Ciò ti consente di consolidare la tua familiarità con i concetti di base di TPUEstimator. Quando il modello è in esecuzione, aggiungi gradualmente più funzionalità.

Guarda i tutorial per un set di modelli di esempio e le istruzioni per eseguirli con Cloud TPU. Sono disponibili modelli aggiuntivi su GitHub.

Per convertire il codice dalla classe tf.estimator.Estimator in modo da utilizzare tf.contrib.tpu.TPUEstimator, modifica quanto segue:

Cambia tf.estimator.RunConfig in tf.contrib.tpu.RunConfig.
Imposta TPUConfig (parte di tf.contrib.tpu.RunConfig) per specificare iterations_per_loop. iterations_per_loop è il numero di iterazioni da eseguire sulla Cloud TPU per una chiamata session.run (per loop di addestramento).

Prima di tornare all'host, Cloud TPU esegue un numero specifico di iterazioni del loop di addestramento. Nessun checkpoint o riepilogo viene salvato finché non vengono eseguite tutte le iterazioni di Cloud TPU.

In model_fn, utilizza tf.contrib.tpu.CrossShardOptimizer per aggregare il tuo strumento di ottimizzazione. Ad esempio:

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

Cambia tf.estimator.Estimator in tf.contrib.tpu.TPUEstimator.

L'impostazione predefinita di RunConfig salva i riepiloghi per TensorBoard ogni 100 passi e scrive i punti di controllo ogni 10 minuti.

Domande frequenti

Perché è richiesto `tf.data` per la pipeline di input?

Ci sono due motivi:

Il codice dell'applicazione viene eseguito sul client mentre il calcolo della TPU viene eseguito sul worker. Per un buon rendimento, le operazioni di pipeline di input devono essere eseguite sul worker. tf.data esegue le operazioni sul worker.
Per ammortizzare il costo del lancio di TPU, il passaggio dell'addestramento del modello viene aggregato in un tf.while_loop, in cui un Session.run esegue molte iterazioni per un singolo ciclo di addestramento. Al momento solo tf.data può essere aggregato in un tf.while_loop.

Come posso configurare il rendimento dell'addestramento di un modello?

Puoi migliorare le prestazioni dell'addestramento del modello utilizzando il profiler fornito per TensorBoard.