Como usar a API TPUEstimator no Cloud TPU

Neste documento, falamos sobre o uso da API TPUEstimator com o Cloud TPU. A TPUEstimator simplifica a execução de modelos em um Cloud TPU ao lidar com vários aspectos específicos de hardware em nível baixo.

Os modelos gravados por meio da TPUEstimator funcionam em CPUs e GPUs, além de dispositivos únicos e pods inteiros de TPU. Geralmente, não são necessárias alterações no código. Ao realizar automaticamente determinadas otimizações para você, a TPUEstimator também garante desempenho máximo com mais facilidade.

API TensorFlow Estimator padrão

Detalhadamente, 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 o modelo para fornecimento.

Além disso, a Estimator tem um comportamento padrão comum nos jobs de treinamento. Isso inclui salvar e restaurar pontos de verificação, criar resumos no TensorBoard etc.

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

Modelo de programação da TPUEstimator

A TPUEstimator encapsula o cálculo, ou model_fn, e o 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 canal 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 fragmentação é 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 os TPUs. O cálculo precisa conter apenas operações compatíveis com o Cloud TPU. As operações do TensorFlow incluem uma lista com todas disponíveis.

Exemplo de treinamento com a TPUEstimator

O código a seguir demonstra o treinamento de um modelo do MNIST usando a 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(
         dataset_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 acima para que você use o Cloud TPU com eficiência.

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:

  1. Na TPUEstimator, o mestre da sessão do TensorFlow não é local. O programa Python cria um único gráfico que é replicado em todos os núcleos no Cloud TPU. Uma configuração comum define o mestre de sessão do TensorFlow como o primeiro trabalhador.

  2. O canal 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.

  3. Os trabalhadores do Cloud TPU funcionam em sincronia, com cada um executando a mesma etapa ao mesmo tempo.

Como converter a TensorFlow Estimator para a TPUEstimator

Recomendamos primeiro a portabilidade de um modelo pequeno e simples e depois o teste do comportamento de ponta a ponta. Isso solidifica seu conhecimento sobre os conceitos básicos da TPUEstimator. Quando o modelo simples for executado, adicione gradualmente mais funcionalidades.

Consulte os tutoriais para ver um conjunto de instruções e modelos de amostra com execução no Cloud TPU. Você encontra mais modelos no GitHub.

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

  • Mude tf.estimator.RunConfig para tf.contrib.tpu.RunConfig.
  • Configure TPUConfig, parte de tf.contrib.tpu.RunConfig, para especificar iterations_per_loop. O iterations_per_loop é o número de iterações a serem executadas no Cloud TPU em uma chamada session.run por ciclo de treinamento.

O Cloud TPU executa um número específico de iterações do ciclo 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 encapsular o otimizador. Por exemplo:

     optimizer = tf.contrib.tpu.CrossShardOptimizer(
          tf.train.GradientDescentOptimizer(learning_rate=learning_rate))
    
  • Mude 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 é necessário no canal de entrada?

Por dois motivos:

  1. O código do aplicativo é executado no cliente enquanto o cálculo do TPU é feito no worker. As operações do canal de entrada precisam estar no trabalhador remoto para você garantir um bom desempenho. Somente tf.data, o conjunto de dados, é compatível.

  2. Para reduzir o custo de lançamento do TPU, a etapa de treinamento do modelo é encapsulada em um tf.while_loop. Assim, um Session.run na verdade executa muitas iterações em um único ciclo de treinamento. No momento, apenas tf.data pode ser encapsulado por um tf.while_loop.

Onde aprendo a executar meu modelo nos Cloud TPUs do Google Cloud Platform?

Siga o tutorial do MNIST e muitos outros neste site para saber mais sobre a execução de modelos na TPUEstimator.

Como crio o perfil de um trabalhador?

Para criar o perfil de um worker, siga as instruções especificadas nas ferramentas do Cloud TPU.

Esta página foi útil? Conte sua opinião sobre:

Enviar comentários sobre…