Exécuter du code TensorFlow sur des tranches de pods Cloud TPU

Ce document explique comment effectuer un calcul à l'aide de TensorFlow sur un pod TPU. Vous allez suivre les étapes suivantes:

  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 exemple de script

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, vous devez y ajouter le rôle Visionneuse TPU. 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 vous créez votre VM TPU.

Configurer votre environnement

  1. Dans Cloud Shell, exécutez la commande suivante pour vous assurer d'exécuter la version actuelle de gcloud:

    $ gcloud components update

    Si vous devez installer gcloud, utilisez la commande suivante:

    $ sudo apt install -y google-cloud-sdk

  2. Créez des variables d'environnement:

    $ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=europe-west4-a
$ export RUNTIME_VERSION=tpu-vm-tf-2.18.0-pod-pjrt
$ export ACCELERATOR_TYPE=v3-32

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=${ZONE} \
  --accelerator-type=${ACCELERATOR_TYPE} \
  --version=${RUNTIME_VERSION}

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 du 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 logicielle de Cloud TPU.

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

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --zone=${ZONE}

Créer et exécuter un exemple de script

  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 sur chaque TensorCore d'une tranche de pod 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=${ZONE}

  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=${ZONE}