Effectuer un calcul sur une VM Cloud TPU à l'aide de TensorFlow
Ce guide de démarrage rapide explique comment créer un Cloud TPU, installer TensorFlow et effectuer un calcul sur un Cloud TPU. Pour un tutoriel plus détaillé vous montrant comment entraîner un modèle sur Cloud TPU, consultez l'un des tutoriels Cloud TPU.
Avant de commencer
Avant de suivre ce guide de démarrage rapide, vous devez créer un compte Google Cloud, installer la Google Cloud CLI et configurer la commande gcloud
.
Pour plus d'informations, consultez la page Configurer un compte et un projet Cloud TPU.
Créer une VM Cloud TPU avec gcloud
Lancez un Cloud TPU Compute Engine à l'aide de la commande gcloud
. Pour en savoir plus sur la commande gcloud
, consultez la documentation de référence sur gcloud
.
Lors du lancement de la VM TPU, vous pouvez utiliser la version du logiciel TPU par défaut indiquée dans la commande ou vous reporter aux versions logicielles de Cloud TPU pour spécifier une autre version de logiciel TPU compatible.
$ gcloud compute tpus tpu-vm create tpu-name
--zone=europe-west4-a
--accelerator-type=v3-8
--version=tpu-vm-tf-2.16.1-pjrt
Description des options de commande
tpu-name
- Nom du Cloud TPU à créer.
zone
- Zone dans laquelle vous prévoyez de créer la ressource 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 la section Versions de TPU.
version
- Version du logiciel TPU.
Se connecter à la VM Cloud TPU
Vous devez vous connecter explicitement à votre VM TPU à l'aide de SSH. Si vous n'êtes pas automatiquement connecté, utilisez la commande suivante.
$ gcloud compute tpus tpu-vm ssh tpu-name
--zone europe-west4-a
Exécuter un exemple avec TensorFlow
Une fois connecté à la VM TPU, définissez la variable d'environnement suivante.
(vm)$ export TPU_NAME=local
Lors de la création de votre TPU, si vous définissez le paramètre --version
sur une version se terminant par -pjrt
, définissez les variables d'environnement suivantes pour activer l'environnement d'exécution PJRT:
(vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
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__) @tf.function def add_fn(x,y): z = x + y return z cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver() tf.config.experimental_connect_to_cluster(cluster_resolver) tf.tpu.experimental.initialize_tpu_system(cluster_resolver) strategy = tf.distribute.TPUStrategy(cluster_resolver) 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 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) }
Effectuer un nettoyage
Pour éviter que les ressources utilisées sur cette page soient facturées sur votre compte Google Cloud, procédez comme suit :
Déconnectez-vous de l'instance Compute Engine, si vous ne l'avez pas déjà fait :
(vm)$ exit
Votre invite devrait maintenant être
username@projectname
, indiquant que vous êtes dans Cloud Shell.Supprimez votre Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name
--zone=europe-west4-aVérifiez que les ressources ont été supprimées en exécutant la commande
gcloud compute tpus tpu-vm list
. La suppression peut prendre plusieurs minutes.$ gcloud compute tpus tpu-vm list --zone=europe-west4-a
Étapes suivantes
Pour en savoir plus sur Cloud TPU, consultez les pages suivantes :