Exécuter le code TensorFlow sur des tranches de pod TPU

Ce document explique comment effectuer un calcul simple à l'aide de TensorFlow sur un pod TPU. Procédez comme suit:

  1. Créer une tranche de pod TPU avec le logiciel TensorFlow
  2. Se connecter à la VM TPU à l'aide de SSH
  3. Créer et exécuter un script simple

La VM TPU s'appuie sur des comptes de service pour obtenir les autorisations d'appeler l'API Cloud TPU. Par défaut, votre VM TPU utilise le compte de service Compute Engine par défaut qui inclut toutes les autorisations Cloud TPU nécessaires. Si vous utilisez votre propre compte de service dont vous avez besoin pour ajouter le lecteur TPU à votre compte de service. Pour en savoir plus sur les rôles Google Cloud, consultez la page Comprendre les rôles. Vous pouvez spécifier votre propre compte de service à l'aide de l'option --service-account lorsque à créer votre VM TPU.

Créer une tranche de pod TPU v3-32 avec l'environnement d'exécution TensorFlow

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=europe-west4-a \
  --accelerator-type=v3-32 \
  --version=tpu-vm-tf-2.17.0-pod-pjrt

Description des options de commande

zone
Zone dans laquelle vous prévoyez de créer votre Cloud TPU.
accelerator-type
Le type d'accélérateur spécifie la version et la taille de la ressource Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
version
Version logicielle de Cloud TPU.

Se connecter à votre VM Cloud TPU à l'aide de SSH

$ gcloud compute tpus tpu-vm ssh tpu-name \
      --zone europe-west4-a

Créer et exécuter un script de calcul simple

  1. Définissez les variables d'environnement suivantes.

    (vm)$ export TPU_NAME=tpu-name
(vm)$ export TPU_LOAD_LIBRARY=0

  2. Créez un fichier nommé tpu-test.py dans le répertoire actuel et collez-y le script suivant.

    import tensorflow as tf
print("Tensorflow version " + tf.__version__)

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
print('Running on TPU ', cluster_resolver.cluster_spec().as_dict()['worker'])

tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)

@tf.function
def add_fn(x,y):
z = x + y
return z

x = tf.constant(1.)
y = tf.constant(1.)
z = strategy.run(add_fn, args=(x,y))
print(z)

  3. Exécutez ce script avec la commande suivante :

    (vm)$ python3 tpu-test.py

    Ce script effectue un calcul simple sur chaque TensorCore d'un TPU. Le résultat ressemblera à ce qui suit :

    PerReplica:{
0: tf.Tensor(2.0, shape=(), dtype=float32),
1: tf.Tensor(2.0, shape=(), dtype=float32),
2: tf.Tensor(2.0, shape=(), dtype=float32),
3: tf.Tensor(2.0, shape=(), dtype=float32),
4: tf.Tensor(2.0, shape=(), dtype=float32),
5: tf.Tensor(2.0, shape=(), dtype=float32),
6: tf.Tensor(2.0, shape=(), dtype=float32),
7: tf.Tensor(2.0, shape=(), dtype=float32),
8: tf.Tensor(2.0, shape=(), dtype=float32),
9: tf.Tensor(2.0, shape=(), dtype=float32),
10: tf.Tensor(2.0, shape=(), dtype=float32),
11: tf.Tensor(2.0, shape=(), dtype=float32),
12: tf.Tensor(2.0, shape=(), dtype=float32),
13: tf.Tensor(2.0, shape=(), dtype=float32),
14: tf.Tensor(2.0, shape=(), dtype=float32),
15: tf.Tensor(2.0, shape=(), dtype=float32),
16: tf.Tensor(2.0, shape=(), dtype=float32),
17: tf.Tensor(2.0, shape=(), dtype=float32),
18: tf.Tensor(2.0, shape=(), dtype=float32),
19: tf.Tensor(2.0, shape=(), dtype=float32),
20: tf.Tensor(2.0, shape=(), dtype=float32),
21: tf.Tensor(2.0, shape=(), dtype=float32),
22: tf.Tensor(2.0, shape=(), dtype=float32),
23: tf.Tensor(2.0, shape=(), dtype=float32),
24: tf.Tensor(2.0, shape=(), dtype=float32),
25: tf.Tensor(2.0, shape=(), dtype=float32),
26: tf.Tensor(2.0, shape=(), dtype=float32),
27: tf.Tensor(2.0, shape=(), dtype=float32),
28: tf.Tensor(2.0, shape=(), dtype=float32),
29: tf.Tensor(2.0, shape=(), dtype=float32),
30: tf.Tensor(2.0, shape=(), dtype=float32),
31: tf.Tensor(2.0, shape=(), dtype=float32)
}

Effectuer un nettoyage

Lorsque vous avez fini d'utiliser votre VM TPU, procédez comme suit pour nettoyer vos ressources.

  1. Déconnectez-vous de Compute Engine:

    (vm)$ exit

  2. Supprimez votre Cloud TPU.

    $ gcloud compute tpus tpu-vm delete tpu-name \
  --zone europe-west4-a

  3. Vérifiez que les ressources ont bien été supprimées en exécutant la commande suivante. Assurez-vous que votre TPU n'est plus répertorié. La suppression peut prendre plusieurs minutes.

    $ gcloud compute tpus tpu-vm list \
  --zone europe-west4-a