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:

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

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

  1. Déconnectez-vous de Compute Engine:

    (vm)$ exit
    
  2. Supprimez votre Cloud TPU.

    $ gcloud compute tpus tpu-vm delete tpu-name \
      --zone europe-west4-a
    
  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 europe-west4-a