Effectuer un calcul sur une VM Cloud TPU à l'aide de TensorFlow

Ce guide de démarrage rapide vous explique comment créer une ressource Cloud TPU, installer TensorFlow et effectuer un calcul simple sur une ressource 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 Platform, 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 ou un nœud Cloud TPU avec gcloud

Lancez un Cloud TPU Compute Engine à l'aide de la commande gcloud. La commande à utiliser varie selon que vous utilisez une VM TPU ou un nœud TPU. Pour en savoir plus sur les deux architectures de VM, consultez la page Architecture du système. Pour en savoir plus sur la commande gcloud, consultez la documentation de référence sur gcloud.

VM TPU

$ 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

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 Cloud TPU.

Nœud TPU

$ gcloud compute tpus execution-groups create \
--name=tpu-name \
--zone=europe-west4-a \
--disk-size=300 \
--machine-type=n1-standard-16 \
--tf-version=2.12.0 \

Description des options de commande

project
Votre ID de projet Google Cloud
name
Nom du Cloud TPU à créer.
zone
Zone dans laquelle vous prévoyez de créer la ressource Cloud TPU.
disk-size
Taille du disque dur en Go de la VM créée par la commande gcloud.
machine-type
Type de machine de la VM Compute Engine à créer.
tf-version
Version de TensorFlow gcloud installée sur la VM. Consultez
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.

Se connecter à la VM Cloud TPU

Lorsque vous utilisez des VM TPU, vous devez vous connecter explicitement à votre VM TPU à l'aide de SSH. Lorsque vous utilisez des nœuds TPU, vous devriez être automatiquement connecté en SSH à votre VM Compute Engine. Si vous n'êtes pas automatiquement connecté, utilisez la commande suivante.

VM TPU

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

Nœud TPU

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

Exécuter un exemple simple avec TensorFlow

VM TPU

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)

Nœud TPU

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

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='your-tpu-name')  # TPU detection
print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

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

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

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.

    VM TPU

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

    Nœud TPU

    $ gcloud compute tpus execution-groups 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.

    VM TPU

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

    Nœud TPU

    $ gcloud compute tpus execution-groups list --zone=europe-west4-a

Étapes suivantes

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