Exécuter un calcul sur une VM Cloud TPU à l'aide de TensorFlow

Ce guide de démarrage rapide explique comment créer un Cloud TPU et exécuter un calcul sur un Cloud TPU à l'aide de TensorFlow. Pour accéder à un tutoriel plus détaillé illustrant l'entraînement d'un modèle sur un 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 Google Cloud CLI et configurer la commande gcloud. Pour en savoir plus, consultez la page Configurer un compte et un projet Cloud TPU.

Créer une VM Cloud TPU avec gcloud

Créez une ressource Cloud TPU à l'aide de la commande gcloud.

  $ gcloud compute tpus tpu-vm create tpu-name 
--zone=europe-west4-a
--accelerator-type=v3-8
--version=tpu-vm-tf-2.17.0-pjrt
--project=your-gcp-project-name

Description des options de commande

tpu-name
Nom du Cloud TPU à créer.
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 du runtime TPU. TensorFlow est préinstallé dans la version utilisée dans ce guide de démarrage rapide.
project
Nom du projet Google Cloud CLI dans lequel vous créez votre Cloud TPU.

Pour en savoir plus sur la commande gcloud, consultez la documentation de référence sur gcloud.

Se connecter à la VM Cloud TPU

Connectez-vous à votre VM TPU à l'aide de SSH:

  $ gcloud compute tpus tpu-vm ssh tpu-name 
--zone europe-west4-a
--project=your-gcp-project-name

Vérifier que TensorFlow peut accéder aux TPU

  1. Créez un fichier nommé tpu-count.py dans le répertoire actuel, puis copiez et collez-y le script suivant.

    import tensorflow as tf
    
    print(f'TensorFlow can access {len(tf.config.list_logical_devices('TPU'))} TPU cores')
    
  2. Exécutez le script :

    (vm)$ python3 tpu-count.py

    La sortie du script indique le nombre de cœurs de TPU disponibles pour la VM TPU :

    TensorFlow can access 8 TPU cores
    

Exécuter un calcul de base à l'aide de TensorFlow

Une fois connecté à la VM TPU, définissez la variable d'environnement suivante.

  (vm)$ export TPU_NAME=local
  

Lorsque vous créez 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 :

  1. 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.

  2. Supprimez votre Cloud TPU.

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

  3. Vé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
      

Étape suivante

Pour en savoir plus sur Cloud TPU, consultez les pages suivantes :