API TPUEstimator no Cloud TPU

Neste documento, você verá o uso da API Qwiklabs com o Cloud TPU. A NetzDG simplifica a execução de modelos em um Cloud TPU ao lidar com vários detalhes específicos de hardware de baixo nível.

Modelos escritos usando a Qwiklabs funcionam em CPUs, GPUs, dispositivos de TPU única e pods de TPU, geralmente sem alterações de código. Ao realizar automaticamente determinadas otimizações para você, a TPUEstimator também garante desempenho máximo com mais facilidade.

Para saber como as cargas de trabalho de machine learning funcionam em hardware de TPU em geral, leia a documentação sobre Arquitetura do sistema.

API TensorFlow Estimator padrão

De modo geral, a API TensorFlow Estimator padrão fornece:

Estimator.train(): treina um modelo em uma determinada entrada em um número fixo de etapas;
Estimator.evaluate(): avalia o modelo em um conjunto de teste;
Estimator.predict(): executa a inferência usando o modelo treinado;
Estimator.export_savedmodel(): exporta seu modelo para veiculação.

Além disso, Estimator inclui o comportamento padrão comum em jobs de treinamento, como salvar e restaurar checkpoints, criar resumos para o TensorBoard e assim por diante.

Estimator requer que você grave um model_fn e um input_fn que correspondam ao modelo e às partes da entrada do gráfico do TensorFlow.

Modelo de programação da TPUEstimator

TPUEstimator encapsula a computação (model_fn) e a distribui para todos os núcleos disponíveis do Cloud TPU. É necessário ajustar a taxa de aprendizado de acordo com o tamanho do lote.

A função input_fn modela o pipeline de entrada em execução na CPU do host remoto. Use tf.data para programar as operações de entrada como descrito no guia do programador. Cada chamada lida com a entrada do lote global em um dispositivo. O tamanho do lote de fragmento é recuperado de params['batch_size']. Dica profissional: para garantir o melhor desempenho, retorne um conjunto de dados em vez de tensores.
A função model_fn modela o cálculo que é replicado e distribuído para as TPUs. O cálculo precisa conter apenas operações compatíveis com o Cloud TPU. Operações do TensorFlow inclui a lista de operações disponíveis.

Exemplo de treinamento com a TPUEstimator

O código a seguir demonstra o treinamento MNIST com 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)

Na seção a seguir, abordamos os novos conceitos apresentados no exemplo anterior, para ajudar você a usar o Cloud TPU de maneira eficaz.

Conceitos da TPUEstimator

A TPUEstimator usa uma abordagem de replicação em gráfico para executar programas do TensorFlow. Essa replicação é de sessão única e diferente do treinamento de replicação entre gráficos ou várias sessões, geralmente usado no TensorFlow distribuído. As principais diferenças incluem:

No TPUEstimator, o iniciador de sessão do TensorFlow não é local. O programa em Python cria um único gráfico que é replicado em todos os núcleos na Cloud TPU. Uma configuração comum define o iniciador de sessão do TensorFlow como o primeiro worker.
O pipeline de entrada é colocado em hosts remotos, em vez de locais, para garantir que os exemplos de treinamento sejam alimentados no Cloud TPU o mais rápido possível. Um conjunto de dados (tf.data) é obrigatório.
Os workers do Cloud TPU funcionam em sincronia, com cada um executando a mesma etapa ao mesmo tempo.

Como converter a TensorFlow Estimator em TPUEstimator

Recomendamos que primeiro você transfira um modelo pequeno e teste o comportamento dele. Isso solidifica seu conhecimento sobre os conceitos básicos da TPUEstimator. Quando o modelo for executado, adicione mais funcionalidades gradualmente.

Consulte os tutoriais para um conjunto de modelos de amostra e instruções para executá-los com o Cloud TPU. Há modelos adicionais disponíveis no GitHub.

Para converter seu código da classe tf.estimator.Estimator para usar tf.contrib.tpu.TPUEstimator, altere o seguinte:

Altere tf.estimator.RunConfig para tf.contrib.tpu.RunConfig.
Defina TPUConfig (parte de tf.contrib.tpu.RunConfig) para especificar iterations_per_loop. iterations_per_loop é o número de iterações a serem executadas no Cloud TPU para uma chamada session.run (por loop de treinamento).

O Cloud TPU executa um número específico de iterações do loop de treinamento antes de retornar ao host. Nenhum ponto de verificação ou resumo é salvo até que todas as iterações do Cloud TPU sejam executadas.

Em model_fn, use tf.contrib.tpu.CrossShardOptimizer para ajustar seu otimizador. Exemplo:

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

Altere tf.estimator.Estimator para tf.contrib.tpu.TPUEstimator.

O RunConfig padrão salva resumos do TensorBoard a cada 100 etapas e grava pontos de verificação a cada 10 minutos.

Perguntas frequentes

Por que `tf.data` é obrigatório para o pipeline de entrada?

Por dois motivos:

O código do aplicativo é executado no cliente enquanto o cálculo da TPU é feito em worker. As operações de pipeline de entrada precisam ser executadas no worker para um bom desempenho. tf.data executa as operações no worker.
Para reduzir o custo de inicialização da TPU, o passo de treinamento do modelo é encapsulado em um tf.while_loop, em que um Session.run executa muitas iterações em um único loop de treinamento. Atualmente, apenas tf.data pode ser unido a um tf.while_loop.

Como crio o perfil do desempenho do treinamento de modelo?

Você pode criar perfis de desempenho de treinamento de modelo usando o criador de perfil fornecido para o TensorBoard.