Use TPUs

Ver no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

Antes de executar esse notebook Colab, certifique-se de que seu acelerador de hardware é uma TPU, verificando suas configurações de notebook: Runtime> Alterar tipo runtime> acelerador de hardware> TPU.

Configurar

import tensorflow as tf

import os
import tensorflow_datasets as tfds

Inicialização TPU

TPUs são normalmente trabalhadores do 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 se conectar ao cluster remoto e inicializar as TPUs. Note que o tpu argumento para tf.distribute.cluster_resolver.TPUClusterResolver é um endereço especial só para Colab. Se você estiver executando seu código no Google Compute Engine (GCE), deverá passar o nome do seu Cloud TPU.

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'))
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.50:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.50:8470
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
All devices:  [LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:5', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:3', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:0', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:1', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:2', device_type='TPU')]

Colocação manual do dispositivo

Depois que a TPU for inicializada, você pode usar o posicionamento manual do dispositivo para colocar o cálculo em um único dispositivo TPU:

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)
c device:  /job:worker/replica:0/task:0/device:TPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

Estratégias de distribuição

Normalmente, você executa seu modelo em várias TPUs de maneira paralela aos dados. Para distribuir seu modelo em vários TPUs (ou outros aceleradores), o TensorFlow oferece várias estratégias de distribuição. Você pode substituir sua estratégia de distribuição e o modelo será executado em qualquer dispositivo (TPU). Verifique a guia de estratégia de distribuição para mais informações.

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

strategy = tf.distribute.TPUStrategy(resolver)
INFO:tensorflow:Found TPU system:
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

Para replicar um cálculo para que ele possa ser executado em todos os núcleos de TPU, você pode passá-lo para o strategy.run API. Abaixo está um exemplo que mostra todos os núcleos que recebe as mesmas entradas (a, b) e da realização da multiplicação da matriz em cada núcleo de forma independente. As saídas serão os valores de todas as réplicas.

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

z = strategy.run(matmul_fn, args=(a, b))
print(z)
PerReplica:{
  0: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  1: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  2: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  3: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  4: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  5: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  6: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32),
  7: tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)
}

Classificação em TPUs

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

Defina um modelo Keras

Comece com uma definição de um Sequential modelo Keras para classificação de imagens no conjunto de dados MNIST usando Keras. Não é diferente do que você usaria se estivesse treinando em CPUs ou GPUs. Note-se que as necessidades de criação de modelo Keras estar dentro strategy.scope , então as variáveis podem ser criados em cada dispositivo TPU. Outras partes do código não precisam estar dentro do escopo da estratégia.

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

Carregue o conjunto de dados

O uso eficiente da tf.data.Dataset API é crítica quando se utiliza um TPU Nuvem, uma vez que é impossível usar o TPU Nuvem menos que você pode alimentá-los de dados com rapidez suficiente. Você pode aprender mais sobre o conjunto de dados de desempenho na guia desempenho do pipeline de entrada .

Para todos, mas os experimentos mais simples (usando tf.data.Dataset.from_tensor_slices ou outros dados em gráfico), você precisa armazenar todos os arquivos de dados lido pelo conjunto de dados no Google Cloud Storage (GCS) baldes.

Para a maioria dos casos de uso, recomenda-se a converter seus dados no TFRecord formato e usar uma tf.data.TFRecordDataset para lê-lo. Verifique o tutorial TFRecord e tf.Example para obter detalhes sobre como fazer isso. Não é uma exigência difícil e você pode usar outros leitores de conjuntos de dados, tais como tf.data.FixedLengthRecordDataset ou tf.data.TextLineDataset .

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

Independentemente do formato de dados usado, é altamente recomendável que você use arquivos grandes da ordem de 100 MB. Isso é especialmente importante nesta configuração de rede, pois a sobrecarga de abertura de um arquivo é significativamente maior.

Como mostrado no código abaixo, você deve usar os tensorflow_datasets módulo para obter uma cópia dos dados de teste formação MNIST e. Note-se que try_gcs é especificado para usar uma cópia que está disponível em um balde GCS público. Se você não especificar isso, a TPU não poderá acessar os dados baixados.

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 Keras

Você pode treinar seu modelo com Keras fit e compile APIs. Não é específico do TPU nada neste passo de escrever o código como se você estivesse usando GPUs mutliple e uma MirroredStrategy vez do TPUStrategy . Você pode saber mais na formação distribuída com Keras tutorial.

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)
Epoch 1/5
300/300 [==============================] - 19s 35ms/step - loss: 0.1409 - sparse_categorical_accuracy: 0.9575 - val_loss: 0.0432 - val_sparse_categorical_accuracy: 0.9860
Epoch 2/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0335 - sparse_categorical_accuracy: 0.9895 - val_loss: 0.0382 - val_sparse_categorical_accuracy: 0.9879
Epoch 3/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0190 - sparse_categorical_accuracy: 0.9941 - val_loss: 0.0360 - val_sparse_categorical_accuracy: 0.9891
Epoch 4/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0129 - sparse_categorical_accuracy: 0.9958 - val_loss: 0.0360 - val_sparse_categorical_accuracy: 0.9894
Epoch 5/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0095 - sparse_categorical_accuracy: 0.9969 - val_loss: 0.0426 - val_sparse_categorical_accuracy: 0.9899
<tensorflow.python.keras.callbacks.History at 0x7f7dac0fafd0>

Para reduzir Python sobrecarga e maximizar o desempenho do seu TPU, passar na argument- steps_per_execution -para Model.compile . Neste exemplo, aumenta o rendimento em cerca de 50%:

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)
Epoch 1/5
300/300 [==============================] - 12s 40ms/step - loss: 0.1325 - sparse_categorical_accuracy: 0.9579 - val_loss: 0.0445 - val_sparse_categorical_accuracy: 0.9853
Epoch 2/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0336 - sparse_categorical_accuracy: 0.9893 - val_loss: 0.0396 - val_sparse_categorical_accuracy: 0.9879
Epoch 3/5
300/300 [==============================] - 3s 9ms/step - loss: 0.0196 - sparse_categorical_accuracy: 0.9937 - val_loss: 0.0404 - val_sparse_categorical_accuracy: 0.9885
Epoch 4/5
300/300 [==============================] - 3s 9ms/step - loss: 0.0113 - sparse_categorical_accuracy: 0.9963 - val_loss: 0.0532 - val_sparse_categorical_accuracy: 0.9872
Epoch 5/5
300/300 [==============================] - 3s 9ms/step - loss: 0.0105 - sparse_categorical_accuracy: 0.9964 - val_loss: 0.0608 - val_sparse_categorical_accuracy: 0.9853
<tensorflow.python.keras.callbacks.History at 0x7f7be03514a8>

Treine o modelo usando um loop de treinamento personalizado

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

Primeiro, crie o modelo, conjuntos de dados e tf.functions:

# Create the model, optimizer and metrics inside the 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 datasets 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)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    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),))
WARNING:tensorflow:From <ipython-input-1-5625c2a14441>:15: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
rename to distribute_datasets_from_function
WARNING:tensorflow:From <ipython-input-1-5625c2a14441>:15: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
rename to distribute_datasets_from_function

Em seguida, execute o loop de treinamento:

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: {}, 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()
Epoch: 0/5
Current step: 300, training loss: 0.1448, accuracy: 95.53%
Epoch: 1/5
Current step: 600, training loss: 0.0332, accuracy: 98.96%
Epoch: 2/5
Current step: 900, training loss: 0.0203, accuracy: 99.36%
Epoch: 3/5
Current step: 1200, training loss: 0.0132, accuracy: 99.59%
Epoch: 4/5
Current step: 1500, training loss: 0.0101, accuracy: 99.62%

Melhorar o desempenho com múltiplos passos dentro tf.function

Você pode melhorar o desempenho, executando várias etapas dentro de um tf.function . Isto é conseguido envolvendo a strategy.run chamada com uma tf.range dentro tf.function , e autógrafo irá convertê-lo para um tf.while_loop no trabalhador TPU.

Apesar do desempenho melhorado, há vantagens e desvantagens com este método em comparação com a execução de um passo único dentro tf.function . A execução de vários passos em uma tf.function é menos flexível, você não pode executar as coisas ansiosamente ou código Python arbitrário dentro das etapas.

@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)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits, from_logits=True)
      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
    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: {}, accuracy: {}%'.format(
      optimizer.iterations.numpy(),
      round(float(training_loss.result()), 4),
      round(float(training_accuracy.result()) * 100, 2)))
Current step: 1800, training loss: 0.0069, accuracy: 99.77%

Próximos passos