Questa pagina è stata tradotta dall'API Cloud Translation.
Switch to English

Usa le TPU

Visualizza su TensorFlow.org Esegui in Google Colab Visualizza sorgente su GitHub Scarica notebook

Il supporto sperimentale per Cloud TPU è attualmente disponibile per Keras e Google Colab. Prima di eseguire questo notebook Colab, assicurati che il tuo acceleratore hardware sia un TPU controllando le impostazioni del tuo notebook: Runtime> Cambia tipo di runtime> Acceleratore hardware> TPU.

Impostare

import tensorflow as tf

import os
import tensorflow_datasets as tfds

Inizializzazione TPU

Le TPU sono in genere sui lavoratori Cloud TPU che sono diversi dal processo locale che esegue il programma Python dell'utente. Pertanto, è necessario eseguire alcune operazioni di inizializzazione per connettersi al cluster remoto e inizializzare le TPU. Nota che l'argomento tpu per TPUClusterResolver è un indirizzo speciale solo per Colab. Nel caso in cui tu stia utilizzando Google Compute Engine (GCE), dovresti invece passare il nome del tuo CloudTPU.

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
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.2:8470

INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2: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: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'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:3', device_type='TPU')]

Posizionamento manuale del dispositivo

Dopo che il TPU è stato inizializzato, è possibile utilizzare il posizionamento manuale del dispositivo per posizionare il calcolo su un singolo 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)

Strategie di distribuzione

La maggior parte delle volte gli utenti desiderano eseguire il modello su più TPU in modo parallelo di dati. Una strategia di distribuzione è un'astrazione che può essere utilizzata per guidare modelli su CPU, GPU o TPU. Sostituisci semplicemente la strategia di distribuzione e il modello verrà eseguito sul dispositivo specificato. Vedere la guida alla strategia di distribuzione per ulteriori informazioni.

Innanzitutto, crea l'oggetto TPUStrategy .

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:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_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)

Per replicare un calcolo in modo che possa essere eseguito in tutti i core TPU, puoi semplicemente passarlo all'API strategy.run . Di seguito è riportato un esempio in cui tutti i core otterranno gli stessi input (a, b) e eseguiranno il matmul su ciascun core in modo indipendente. Gli output saranno i valori di tutte le repliche.

@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)
}

Classificazione su TPU

Poiché abbiamo appreso i concetti di base, è tempo di guardare a un esempio più concreto. Questa guida mostra come utilizzare la strategia di distribuzione tf.distribute.experimental.TPUStrategy per guidare una Cloud TPU e addestrare un modello Keras.

Definisci un modello Keras

Di seguito la definizione del modello MNIST che utilizza Keras, invariato rispetto a quello che useresti su CPU o GPU. Nota che la creazione del modello Keras deve essere all'interno di strategy.scope , quindi le variabili possono essere create su ogni dispositivo TPU. Altre parti del codice non sono necessarie per essere all'interno dell'ambito della strategia.

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)])

Set di dati di input

L'uso efficiente dell'API tf.data.Dataset è fondamentale quando si utilizza una Cloud TPU, poiché è impossibile utilizzare le Cloud TPU a meno che non sia possibile fornire loro i dati abbastanza rapidamente. Vedere Guida alle prestazioni della pipeline di input per i dettagli sulle prestazioni del set di dati.

Per tutti tranne la sperimentazione più semplice (utilizzando tf.data.Dataset.from_tensor_slices o altri dati nel grafico) dovrai memorizzare tutti i file di dati letti dal set di dati nei bucket di Google Cloud Storage (GCS).

Per la maggior parte dei casi d'uso, si consiglia di convertire i dati in formato TFRecord e utilizzare un tf.data.TFRecordDataset per leggerli. Vedi TFRecord e il tutorial tf.Example per i dettagli su come farlo. Questo, tuttavia, non è un requisito fondamentale e puoi utilizzare altri lettori di set di dati ( FixedLengthRecordDataset o TextLineDataset ) se preferisci.

tf.data.Dataset.cache dati di piccole dimensioni possono essere caricati interamente in memoria utilizzando tf.data.Dataset.cache .

Indipendentemente dal formato dei dati utilizzato, si consiglia vivamente di utilizzare file di grandi dimensioni, nell'ordine di 100 MB. Ciò è particolarmente importante in questa impostazione di rete poiché il sovraccarico di apertura di un file è notevolmente maggiore.

Qui dovresti usare il modulo tensorflow_datasets per ottenere una copia dei dati di addestramento MNIST. Tieni presente che try_gcs è specificato per utilizzare una copia disponibile in un bucket GCS pubblico. Se non lo specifichi, il TPU non sarà in grado di accedere ai dati scaricati.

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)

  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 to have a
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so users 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

Addestra un modello utilizzando le API di alto livello di Keras

Puoi addestrare un modello semplicemente con le API di adattamento / compilazione di Keras. Niente qui è specifico per TPU, scriveresti lo stesso codice di seguito se avessi più GPU e utilizzassi una TPUStrategy MirroredStrategy anziché una TPUStrategy . Per saperne di più, guarda il tutorial Formazione distribuita con Keras .

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
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Iterator.get_next_as_optional()` instead.

Warning:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Iterator.get_next_as_optional()` instead.

  1/300 [..............................] - ETA: 20:25 - loss: 2.3062 - sparse_categorical_accuracy: 0.0500WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0015s vs `on_train_batch_end` time: 0.0237s). Check your callbacks.

Warning:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0015s vs `on_train_batch_end` time: 0.0237s). Check your callbacks.

300/300 [==============================] - ETA: 0s - loss: 0.1450 - sparse_categorical_accuracy: 0.9561WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0013s vs `on_test_batch_end` time: 0.0097s). Check your callbacks.

Warning:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0013s vs `on_test_batch_end` time: 0.0097s). Check your callbacks.

300/300 [==============================] - 15s 50ms/step - loss: 0.1450 - sparse_categorical_accuracy: 0.9561 - val_loss: 0.0456 - val_sparse_categorical_accuracy: 0.9852
Epoch 2/5
300/300 [==============================] - 8s 28ms/step - loss: 0.0357 - sparse_categorical_accuracy: 0.9890 - val_loss: 0.0372 - val_sparse_categorical_accuracy: 0.9884
Epoch 3/5
300/300 [==============================] - 9s 28ms/step - loss: 0.0193 - sparse_categorical_accuracy: 0.9940 - val_loss: 0.0557 - val_sparse_categorical_accuracy: 0.9835
Epoch 4/5
300/300 [==============================] - 9s 29ms/step - loss: 0.0141 - sparse_categorical_accuracy: 0.9954 - val_loss: 0.0405 - val_sparse_categorical_accuracy: 0.9883
Epoch 5/5
300/300 [==============================] - 9s 29ms/step - loss: 0.0092 - sparse_categorical_accuracy: 0.9967 - val_loss: 0.0428 - val_sparse_categorical_accuracy: 0.9887

<tensorflow.python.keras.callbacks.History at 0x7f480c793240>

Per ridurre il sovraccarico di Python e massimizzare le prestazioni della tua TPU, prova l'argomento experiment_steps_per_execution experimental_steps_per_execution su Model.compile . Qui aumenta la produttività di circa il 50%:

with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                # Anything between 2 and `steps_per_epoch` could help here.
                experimental_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 [==============================] - 16s 52ms/step - loss: 0.1337 - sparse_categorical_accuracy: 0.9583 - val_loss: 0.0521 - val_sparse_categorical_accuracy: 0.9840
Epoch 2/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0331 - sparse_categorical_accuracy: 0.9898 - val_loss: 0.0360 - val_sparse_categorical_accuracy: 0.9884
Epoch 3/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0188 - sparse_categorical_accuracy: 0.9939 - val_loss: 0.0405 - val_sparse_categorical_accuracy: 0.9887
Epoch 4/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0117 - sparse_categorical_accuracy: 0.9961 - val_loss: 0.0808 - val_sparse_categorical_accuracy: 0.9820
Epoch 5/5
300/300 [==============================] - 5s 16ms/step - loss: 0.0119 - sparse_categorical_accuracy: 0.9962 - val_loss: 0.0488 - val_sparse_categorical_accuracy: 0.9862

<tensorflow.python.keras.callbacks.History at 0x7f47ac791240>

Addestra un modello utilizzando il ciclo di addestramento personalizzato.

Puoi anche creare e addestrare i tuoi modelli utilizzando direttamente le API tf.function e tf.distribute . strategy.experimental_distribute_datasets_from_function API strategy.experimental_distribute_datasets_from_function viene utilizzata per distribuire il set di dati data una funzione del set di dati. Si noti che la dimensione del batch passata nel set di dati sarà per la dimensione del batch di replica invece della dimensione del batch globale in questo caso. Per saperne di più, consulta il tutorial Formazione personalizzata con tf.distribute.Strategy .

Innanzitutto, crea il modello, i set di dati e le funzioni tf.

# Create the model, optimizer and metrics inside 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),))

Quindi esegui il ciclo di allenamento.

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.1335, accuracy: 95.8%
Epoch: 1/5
Current step: 600, training loss: 0.0344, accuracy: 98.93%
Epoch: 2/5
Current step: 900, training loss: 0.0196, accuracy: 99.35%
Epoch: 3/5
Current step: 1200, training loss: 0.0119, accuracy: 99.61%
Epoch: 4/5
Current step: 1500, training loss: 0.0102, accuracy: 99.65%

Miglioramento delle prestazioni di più passaggi all'interno della funzione tf.function

Le prestazioni possono essere migliorate eseguendo più passaggi all'interno di una funzione tf.function .. Ciò si ottiene avvolgendo la chiamata strategy.run con un tf.range all'interno di tf.function , AutoGraph lo convertirà in un tf.while_loop sul worker TPU.

Sebbene con prestazioni migliori, ci sono compromessi rispetto a un singolo passaggio all'interno di tf.function . L'esecuzione di più passaggi in una funzione tf.function è meno flessibile, non è possibile eseguire le cose con entusiasmo o codice Python arbitrario all'interno dei passaggi.

@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.0087, accuracy: 99.72%

Prossimi passi