Exécuter du code TensorFlow sur des tranches de pod TPU
Ce document explique comment effectuer un calcul simple à l'aide de TensorFlow sur un pod TPU. Voici les étapes à suivre:
- Créer une tranche de pod TPU avec le logiciel TensorFlow
- Se connecter à la VM TPU à l'aide de SSH
- Créer et exécuter un script simple
La VM TPU s'appuie sur un compte de service pour les autorisations nécessaires pour 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, vous devez ajouter le rô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
lors de la création de 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.16.1-pod-pjrt
Description des options de commande
zone
- Zone dans laquelle vous prévoyez de créer la ressource Cloud TPU.
accelerator-type
- Type de Cloud TPU à créer.
version
- Version du logiciel Cloud TPU.
Se connecter à la 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
Définissez les variables d'environnement suivantes.
(vm)$ export TPU_NAME=tpu-name (vm)$ export TPU_LOAD_LIBRARY=0
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)
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.
Déconnectez-vous de Compute Engine:
(vm)$ exit
Supprimez votre Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone europe-west4-a
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