Utiliser des TPU

Voir sur TensorFlow.org Exécuter dans Google Colab Voir la source sur GitHub Télécharger le cahier

Avant d'exécuter ce bloc - notes Colab, assurez - vous que votre accélérateur matériel est un TPU en vérifiant vos paramètres de bloc - notes: Runtime> Modifier le type d'exécution> accélérateur matériel> TPU.

Installer

import tensorflow as tf

import os
import tensorflow_datasets as tfds
/home/kbuilder/.local/lib/python3.6/site-packages/requests/__init__.py:104: RequestsDependencyWarning: urllib3 (1.26.6) or chardet (2.3.0)/charset_normalizer (2.0.4) doesn't match a supported version!
  RequestsDependencyWarning)

Initialisation TPU

Les TPU sont généralement des nœuds de calcul Cloud TPU, qui sont différents du processus local exécutant le programme Python de l'utilisateur. Ainsi, vous devez effectuer un travail d'initialisation pour vous connecter au cluster distant et initialiser les TPU. Notez que le tpu argument tf.distribute.cluster_resolver.TPUClusterResolver est une adresse spéciale pour Colab. Si vous exécutez votre code sur Google Compute Engine (GCE), vous devez plutôt transmettre le nom de votre 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:Clearing out eager caches
INFO:tensorflow:Clearing out eager caches
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:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
All devices:  [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'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', 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:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU')]

Placement manuel de l'appareil

Une fois le TPU initialisé, vous pouvez utiliser le placement manuel de l'appareil pour placer le calcul sur un seul appareil 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)

Stratégies de distribution

Généralement, vous exécutez votre modèle sur plusieurs TPU de manière parallèle aux données. Pour distribuer votre modèle sur plusieurs TPU (ou autres accélérateurs), TensorFlow propose plusieurs stratégies de distribution. Vous pouvez remplacer votre stratégie de distribution et le modèle s'exécutera sur n'importe quel périphérique (TPU) donné. Consultez le guide de stratégie de distribution pour plus d' informations.

Pour démontrer cela, créez un tf.distribute.TPUStrategy objet:

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)

Pour reproduire un calcul il peut fonctionner dans tous les cœurs de TPU, vous pouvez passer dans l' strategy.run API. Ci - dessous est un exemple qui montre tous les cœurs recevant les mêmes entrées (a, b) et effectuer une multiplication matricielle sur chaque noyau de façon indépendante. Les sorties seront les valeurs de toutes les répliques.

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

Classification sur les TPU

Après avoir couvert les concepts de base, considérons un exemple plus concret. Cette section montre comment utiliser de la distribution tf.distribute.TPUStrategy -pour former un modèle Keras sur un nuage TPU.

Définir un modèle Keras

Commencez par une définition d'un Sequential modèle Keras pour la classification d'image sur l'ensemble de données en utilisant MNIST Keras. Ce n'est pas différent de ce que vous utiliseriez si vous vous entraîniez sur des CPU ou des GPU. Notez que les besoins de création de modèle Keras être à l' intérieur strategy.scope , de sorte que les variables peuvent être créées sur chaque appareil de TPU. Il n'est pas nécessaire que d'autres parties du code entrent dans le champ d'application de la stratégie.

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

Charger le jeu de données

L' utilisation efficace de l' tf.data.Dataset API est critique lorsque vous utilisez un nuage TPU, car il est impossible d'utiliser le Cloud PUT à moins que vous pouvez les nourrir des données assez rapidement. Vous pouvez en savoir plus sur les performances ensemble de données dans le guide de la performance du pipeline d' entrée .

Pour tous , mais les expériences les plus simples ( en utilisant tf.data.Dataset.from_tensor_slices ou autre en graphique des données), vous devez stocker tous les fichiers de données lus par les seaux Dataset dans Google Cloud Storage (GCS).

Pour la plupart des cas d'utilisation, il est recommandé de convertir vos données dans le TFRecord le format et utiliser un tf.data.TFRecordDataset pour le lire. Vérifiez le tutoriel TFRecord et tf.Example pour plus de détails sur la façon de le faire. Il n'est pas une exigence difficile et vous pouvez utiliser d' autres lecteurs de jeux de données, tels que tf.data.FixedLengthRecordDataset ou tf.data.TextLineDataset .

Vous pouvez charger des jeux de données entiers petits dans la mémoire à l' aide tf.data.Dataset.cache .

Quel que soit le format de données utilisé, il est fortement recommandé d'utiliser des fichiers volumineux de l'ordre de 100 Mo. Ceci est particulièrement important dans ce cadre en réseau, car la surcharge d'ouverture d'un fichier est considérablement plus élevée.

Comme le montre le code ci - dessous, vous devez utiliser le tensorflow_datasets module pour obtenir une copie de la formation MNIST et des données de test. Notez que try_gcs est spécifié à utiliser une copie qui est disponible dans un seau de GCS public. Si vous ne le spécifiez pas, le TPU ne pourra pas accéder aux données téléchargées.

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

Entraîner le modèle à l'aide des API de haut niveau Keras

Vous pouvez former votre modèle avec Keras fit et de compile des API. Il y a dans cette étape-vous spécifique à TPU rien écrire le code comme si vous utilisiez processeurs graphiques et un mutliple MirroredStrategy au lieu du TPUStrategy . Vous pouvez en savoir plus dans la formation distribuée avec Keras tutoriel.

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.1352 - sparse_categorical_accuracy: 0.9583 - val_loss: 0.0404 - val_sparse_categorical_accuracy: 0.9863
Epoch 2/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0346 - sparse_categorical_accuracy: 0.9893 - val_loss: 0.0376 - val_sparse_categorical_accuracy: 0.9882
Epoch 3/5
300/300 [==============================] - 7s 22ms/step - loss: 0.0199 - sparse_categorical_accuracy: 0.9934 - val_loss: 0.0432 - val_sparse_categorical_accuracy: 0.9872
Epoch 4/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0132 - sparse_categorical_accuracy: 0.9957 - val_loss: 0.0475 - val_sparse_categorical_accuracy: 0.9870
Epoch 5/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0100 - sparse_categorical_accuracy: 0.9970 - val_loss: 0.0443 - val_sparse_categorical_accuracy: 0.9877
<keras.callbacks.History at 0x7faf8829f278>

Pour réduire les frais généraux Python et optimiser les performances de votre TPU, passer dans le argument- steps_per_execution -pour Model.compile . Dans cet exemple, il augmente le débit d'environ 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 [==============================] - 13s 43ms/step - loss: 0.1405 - sparse_categorical_accuracy: 0.9571 - val_loss: 0.0438 - val_sparse_categorical_accuracy: 0.9860
Epoch 2/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0337 - sparse_categorical_accuracy: 0.9890 - val_loss: 0.0503 - val_sparse_categorical_accuracy: 0.9847
Epoch 3/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0170 - sparse_categorical_accuracy: 0.9947 - val_loss: 0.0448 - val_sparse_categorical_accuracy: 0.9884
Epoch 4/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0137 - sparse_categorical_accuracy: 0.9956 - val_loss: 0.0362 - val_sparse_categorical_accuracy: 0.9903
Epoch 5/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0090 - sparse_categorical_accuracy: 0.9971 - val_loss: 0.0458 - val_sparse_categorical_accuracy: 0.9880
<keras.callbacks.History at 0x7fae1054f908>

Entraîner le modèle à l'aide d'une boucle d'entraînement personnalisée

Vous pouvez également créer et former votre modèle en utilisant tf.function et tf.distribute API directement. Vous pouvez utiliser le strategy.experimental_distribute_datasets_from_function API pour distribuer l'ensemble de données en fonction du jeu de données. Notez que dans l'exemple ci-dessous, la taille de lot transmise à l'ensemble de données est la taille de lot par réplica au lieu de la taille de lot globale. Pour en savoir plus, consultez la formation personnalisée avec tf.distribute.Strategy tutoriel.

Tout d'abord, créez le modèle, les ensembles de données et les fonctions tf. :

# 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

Ensuite, exécutez la boucle d'entraînement :

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.1352, accuracy: 95.85%
Epoch: 1/5
Current step: 600, training loss: 0.0355, accuracy: 98.86%
Epoch: 2/5
Current step: 900, training loss: 0.0194, accuracy: 99.37%
Epoch: 3/5
Current step: 1200, training loss: 0.0127, accuracy: 99.6%
Epoch: 4/5
Current step: 1500, training loss: 0.0111, accuracy: 99.65%

L' amélioration de la performance avec des étapes multiples à l' intérieur tf.function

Vous pouvez améliorer les performances en exécutant plusieurs étapes dans un tf.function . Ceci est réalisé en enveloppant le strategy.run appel avec un tf.range intérieur tf.function et AutoGraph convertira à un tf.while_loop sur le travailleur de TPU.

En dépit de l'amélioration de la performance, il y a des compromis avec cette méthode par rapport à l' exécution d' une seule étape à l' intérieur tf.function . Exécution de plusieurs étapes dans un tf.function est moins flexible des choses que vous ne pouvez pas exécuter avec impatience ou code Python arbitraire dans les étapes.

@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.0081, accuracy: 99.73%

Prochaines étapes