GitHub Repository: tensorflow/docs-l10n
Path: blob/master/site/pt-br/guide/tpu.ipynb
²⁵¹¹⁵ views

Kernel: Python 3

Copyright 2018 The TensorFlow Authors.

In [ ]:

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Use TPUs

Veja em TensorFlow.org

Executar no Google Colab

Ver fonte em GitHub

Baixar notebook

Este guia demonstra como realizar treinamento básico em Unidades de Processamento de Tensores (TPUs) e TPU Pods, uma coleção de dispositivos de TPU conectados por interfaces de rede de alta velocidade dedicadas, com tf.keras e loops de treinamento personalizados.

TPUs são circuitos integrados de aplicação específica (ASICs) desenvolvidos sob medida pelo Google e usados para acelerar as cargas de trabalho de aprendizado de máquina. Eles estão disponíveis no Google Colab, TPU Research Cloud e Cloud TPU.

Configuração

Antes de executar este notebook do Colab, confira se o acelerador de hardware é um TPU ao verificar as configurações do seu notebook: Runtime > Change runtime type > Hardware accelerator > TPU.

Importe algumas bibliotecas necessárias, incluindo o TensorFlow Datasets:

In [ ]:

import tensorflow as tf

import os
import tensorflow_datasets as tfds

Inicialização da TPU

TPUs são tipicamente workers da Cloud TPU, que são diferentes do processo local que executa o programa Python do usuário. Portanto, você precisa fazer algum trabalho de inicialização para conectar-se ao cluster remoto e inicializar as TPUs. Observe que o argumento tpu recebido por tf.distribute.cluster_resolver.TPUClusterResolver é um endereço especial apenas para o Colab. Se você estiver executando seu código no Google Compute Engine (GCE), deverá passar o nome do Cloud TPU.

Observação: O código de inicialização da TPU deve estar no início do seu programa.

In [ ]:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

Posicionamento manual dos dispositivos

Depois que a TPU for inicializada, você poderá usar o posicionamento manual do dispositivo para colocar a computação num único dispositivo de TPU:

In [ ]:

a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

with tf.device('/TPU:0'):
  c = tf.matmul(a, b)

print("c device: ", c.device)
print(c)

Estratégias de distribuição

Geralmente você executará seu modelo em múltiplas TPUs de maneira paralela aos dados. Para distribuir seu modelo em múltiplas TPUs (assim como em múltiplas GPUs ou em múltiplas máquinas), o TensorFlow oferece a API tf.distribute.Strategy. Você pode substituir sua estratégia de distribuição e o modelo será executado em qualquer dispositivo (TPU). Saiba mais no guia Treinamento distribuído com TensorFlow.

O uso da opção tf.distribute.TPUStrategy implementa treinamento distribuído síncrono. As TPUs fornecem sua própria implementação de operações all-reduce eficientes e outras operações coletivas em múltiplos cores de TPU, que são usados na TPUStrategy.

Para demonstrar isso, crie um objeto tf.distribute.TPUStrategy:

In [ ]:

strategy = tf.distribute.TPUStrategy(resolver)

Para replicar uma computação para que ela possa ser executada em todos os cores da TPU, você pode passá-la para a API Strategy.run. Abaixo está um exemplo que mostra todos os cores recebendo as mesmas entradas (a, b) e realizando multiplicação matricial em cada core de forma independente. As saídas serão os valores obtidos de todas as réplicas.

In [ ]:

@tf.function
def matmul_fn(x, y):
  z = tf.matmul(x, y)
  return z

z = strategy.run(matmul_fn, args=(a, b))
print(z)

Classificação em TPUs

Tendo coberto os conceitos básicos, considere um exemplo mais concreto. Esta seção demonstra como usar a estratégia de distribuição (tf.distribute.TPUStrategy) para treinar um modelo Keras numa Cloud TPU.

Defina um modelo Keras

Comece com uma definição de um modelo Sequential Keras para classificação de imagens no dataset MNIST. Não é diferente do que você usaria se estivesse treinando em CPUs ou GPUs. Observe que a criação do modelo Keras precisa estar dentro do Strategy.scope, para que as variáveis possam ser criadas em cada dispositivo TPU. Outras partes do código não são necessárias para estarem dentro do escopo de Strategy.

In [ ]:

def create_model():
  regularizer = tf.keras.regularizers.L2(1e-5)
  return tf.keras.Sequential(
      [tf.keras.layers.Conv2D(256, 3, input_shape=(28, 28, 1),
                              activation='relu',
                              kernel_regularizer=regularizer),
       tf.keras.layers.Conv2D(256, 3,
                              activation='relu',
                              kernel_regularizer=regularizer),
       tf.keras.layers.Flatten(),
       tf.keras.layers.Dense(256,
                             activation='relu',
                             kernel_regularizer=regularizer),
       tf.keras.layers.Dense(128,
                             activation='relu',
                             kernel_regularizer=regularizer),
       tf.keras.layers.Dense(10,
                             kernel_regularizer=regularizer)])

Este modelo coloca os termos de regularização L2 nos pesos de cada camada, para que o loop de treinamento personalizado abaixo possa mostrar como obtê-los de Model.losses.

Carregue o dataset

O uso eficiente da API tf.data.Dataset é fundamental ao usar uma Cloud TPU. Você pode saber mais sobre o desempenho do dataset no Guia de desempenho do pipeline de entrada.

Se você estiver usando TPU Nodes, precisará armazenar todos os arquivos de dados lidos pelo Dataset do TensorFlow nos buckets do Google Cloud Storage (GCS). Se você estiver usando TPU VMs, poderá armazenar dados onde quiser. Para mais informações sobre TPU Nodes e TPU VMs, consulte a documentação da Arquitetura do Sistema TPU.

Para a maioria dos casos de uso, é recomendado converter seus dados para o formato TFRecord e usar tf.data.TFRecordDataset para lê-los. Verifique o tutorial sobre TFRecord e tf.Example para detalhes sobre como fazer isso. Não é um requisito difícil e você pode usar outros leitores de dataset, como tf.data.FixedLengthRecordDataset ou tf.data.TextLineDataset.

Você pode carregar pequenos datasets inteiros na memória usando tf.data.Dataset.cache.

Independentemente do formato de dados utilizado, é altamente recomendável usar arquivos grandes, da ordem de 100 MB. Isto é especialmente importante neste ambiente de rede, pois o overhead na abertura de um arquivo é significativamente maior.

Conforme mostrado no código abaixo, você deve usar o módulo tfds.load do Tensorflow Datasets para obter uma cópia dos dados de treinamento e teste do MNIST. Observe que try_gcs é especificado para usar uma cópia que está disponível num bucket público do GCS. Se você não especificar isso, a TPU não poderá acessar os dados baixados.

In [ ]:

def get_dataset(batch_size, is_training=True):
  split = 'train' if is_training else 'test'
  dataset, info = tfds.load(name='mnist', split=split, with_info=True,
                            as_supervised=True, try_gcs=True)

  # Normalize the input data.
  def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255.0
    return image, label

  dataset = dataset.map(scale)

  # Only shuffle and repeat the dataset in training. The advantage of having an
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so that you don't need to think about scaling the gradients
  # based on the actual batch size.
  if is_training:
    dataset = dataset.shuffle(10000)
    dataset = dataset.repeat()

  dataset = dataset.batch(batch_size)

  return dataset

Treine o modelo usando APIs de alto nível do Keras

Você pode treinar seu modelo com APIs Keras Model.fit e Model.compile. Não há nada específico de TPU nesta etapa – você escreve o código como se estivesse usando múltiplas GPUs e uma MirroredStrategy em vez de TPUStrategy. Você pode aprender mais no tutorial Treinamento distribuído com Keras.

In [ ]:

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

batch_size = 200
steps_per_epoch = 60000 // batch_size
validation_steps = 10000 // batch_size

train_dataset = get_dataset(batch_size, is_training=True)
test_dataset = get_dataset(batch_size, is_training=False)

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset,
          validation_steps=validation_steps)

Para reduzir o overhead do Python e maximizar o desempenho da sua TPU, passe o argumento steps_per_execution para o Keras Model.compile. Neste exemplo, isto aumenta o rendimento em cerca de 50%:

In [ ]:

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                # Anything between 2 and `steps_per_epoch` could help here.
                steps_per_execution = 50,
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset,
          validation_steps=validation_steps)

Treine o modelo usando um loop de treinamento personalizado

Você também pode criar e treinar seu modelo usando as APIs tf.function e tf.distribute diretamente. Você pode usar a API Strategy.experimental_distribute_datasets_from_function para distribuir o tf.data.Dataset dada uma função de dataset. Observe que no exemplo abaixo o tamanho do lote passado para o Dataset é o tamanho do lote por réplica em vez do tamanho do lote global. Para saber mais, confira o tutorial Treinamento personalizado com tf.distribute.Strategy.

Primeiro, crie o modelo, os datasets e as tf.function:

In [ ]:

# Create the model, optimizer and metrics inside the `tf.distribute.Strategy`
# scope, so that the variables can be mirrored on each device.
with strategy.scope():
  model = create_model()
  optimizer = tf.keras.optimizers.Adam()
  training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)
  training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      'training_accuracy', dtype=tf.float32)

# Calculate per replica batch size, and distribute the `tf.data.Dataset`s
# on each TPU worker.
per_replica_batch_size = batch_size // strategy.num_replicas_in_sync

train_dataset = strategy.experimental_distribute_datasets_from_function(
    lambda _: get_dataset(per_replica_batch_size, is_training=True))

@tf.function
def train_step(iterator):
  """The step function for one training step."""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      per_example_loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(per_example_loss)
      model_losses = model.losses
      if model_losses:
        loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))

    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  strategy.run(step_fn, args=(next(iterator),))

Em seguida, execute o loop de treinamento:

In [ ]:

steps_per_eval = 10000 // batch_size

train_iterator = iter(train_dataset)
for epoch in range(5):
  print('Epoch: {}/5'.format(epoch))

  for step in range(steps_per_epoch):
    train_step(train_iterator)
  print('Current step: {}, training loss: {}, training accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))
  training_loss.reset_states()
  training_accuracy.reset_states()

Melhorando o desempenho com múltiplos passos dentro de `tf.function`

Você pode melhorar o desempenho executando múltiplos passos numa tf.function. Isso é conseguido empacotando a chamada Strategy.run com um tf.range dentro de tf.function, e o AutoGraph irá convertê-la num tf.while_loop no worker TPU. Você pode aprender mais sobre tf.function no guia Melhor desempenho com tf.function.

Apesar do desempenho aprimorado, há vantagens e desvantagens nesse método em comparação com a execução de um único passo dentro de um tf.function. Executar vários passos em uma tf.function é menos flexível – você não pode executar coisas de forma eager ou código Python arbitrário dentro dos passos.

In [ ]:

@tf.function
def train_multiple_steps(iterator, steps):
  """The step function for one training step."""

  def step_fn(inputs):
    """The computation to run on each TPU device."""
    images, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      per_example_loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(per_example_loss)
      model_losses = model.losses
      if model_losses:
        loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
    training_loss.update_state(loss * strategy.num_replicas_in_sync)
    training_accuracy.update_state(labels, logits)

  for _ in tf.range(steps):
    strategy.run(step_fn, args=(next(iterator),))

# Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get
# retraced if the value changes.
train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))

print('Current step: {}, training loss: {}, training accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))

Próximos passos

Para saber mais sobre as Cloud TPUs e como usá-las, consulte:

Google Cloud TPU: a página principal da Google Cloud TPU.
Documentação da Google Cloud TPU documentation: toda a documentação da Google Cloud TPU, que inclui:
- Introdução à Cloud TPU: uma visão geral de como trabalhar com Cloud TPUs.
- Guias de início rápido do Cloud TPU: guias introdutórios de como trabalhar com VMs da Cloud TPU usando o TensorFlow e outros frameworks de machine learning.
Notebooks Colab do Google Cloud TPU: exemplos de treinamento completo.
Guia de desempenho do Google Cloud TPU: melhore ainda mais o desempenho do Cloud TPU ajustando os parâmetros de configuração do Cloud TPU para seu aplicativo
Treinamento distribuído com TensorFlow: como usar estratégias de distribuição, incluindo tf.distribute.TPUStrategy, com exemplos que mostram as práticas recomendadas.
Embeddings de TPU: o TensorFlow inclui suporte especializado para treinamento de embeddings em TPUs por meio de tf.tpu.experimental.embedding. Além disso, o TensorFlow Tecommenders possui tfrs.layers.embedding.TPUEmbedding. Os embeddings fornecem representações eficientes e densas, capturando similaridades e relacionamentos complexos entre características. O suporte a embeddings específicos para TPU do TensorFlow permite treinar embeddings maiores que a memória de um único dispositivo TPU e usar entradas esparsas e irregulares em TPUs.
TPU Research Cloud (TRC): o TRC permite que pesquisadores solicitem acesso a um cluster de mais de 1.000 dispositivos Cloud TPU.

Copyright 2018 The TensorFlow Authors.

Use TPUs

Configuração

Inicialização da TPU

Posicionamento manual dos dispositivos

Estratégias de distribuição

Classificação em TPUs

Defina um modelo Keras

Carregue o dataset

Treine o modelo usando APIs de alto nível do Keras

Treine o modelo usando um loop de treinamento personalizado

Melhorando o desempenho com múltiplos passos dentro de `tf.function`

Próximos passos

Product

Resources

Company

Copyright 2018 The TensorFlow Authors.

Use TPUs

Configuração

Inicialização da TPU

Posicionamento manual dos dispositivos

Estratégias de distribuição

Classificação em TPUs

Defina um modelo Keras

Carregue o dataset

Treine o modelo usando APIs de alto nível do Keras

Treine o modelo usando um loop de treinamento personalizado

Melhorando o desempenho com múltiplos passos dentro de tf.function

Próximos passos

Melhorando o desempenho com múltiplos passos dentro de `tf.function`