Esegui il codice TensorFlow sulle sezioni dei pod di TPU

Questo documento mostra come eseguire un semplice calcolo utilizzando TensorFlow su un pod di TPU. Eseguirai questi passaggi:

  1. Crea una sezione pod di TPU con il software TensorFlow
  2. Connettiti alla VM TPU tramite SSH
  3. Creare ed eseguire uno script semplice

La VM TPU si basa su un account di servizio per le autorizzazioni per chiamare l'API Cloud TPU. Per impostazione predefinita, la VM TPU utilizza l'account di servizio predefinito di Compute Engine, che include tutte le autorizzazioni Cloud TPU necessarie. Se utilizzi il tuo account di servizio, devi aggiungere il ruolo Visualizzatore TPU all'account di servizio. Per ulteriori informazioni sui ruoli di Google Cloud, vedi Informazioni sui ruoli. Puoi specificare il tuo account di servizio utilizzando il flag --service-account durante la creazione della VM TPU.

Crea una sezione pod di TPU v3-32 con runtime TensorFlow

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=europe-west4-a \
  --accelerator-type=v3-32 \
  --version=tpu-vm-tf-2.13.0-pod

Descrizioni dei flag di comando

zone
La zona in cui prevedi di creare il tuo Cloud TPU.
accelerator-type
Il tipo di Cloud TPU da creare.
version
La versione del software Cloud TPU . L'impostazione predefinita è la versione più recente del software TensorFlow.

Connettiti alla tua VM Cloud TPU tramite SSH

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

Crea ed esegui un semplice script di calcolo

  1. Imposta le variabili di ambiente riportate di seguito.

    (vm)$ export TPU_NAME=tpu-name
    (vm)$ export TPU_LOAD_LIBRARY=0
    
  2. Crea un file denominato tpu-test.py nella directory attuale, copialo e incolla al suo interno il seguente script.

    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. Esegui questo script con il seguente comando:

    (vm)$ python3 tpu-test.py

    Questo script esegue un semplice calcolo su ogni TensorCore di una TPU. L'output sarà simile al seguente:

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

Esegui la pulizia

Quando hai finito di utilizzare la VM TPU, segui questi passaggi per ripulire le risorse.

  1. Disconnetti da Compute Engine:

    (vm)$ exit
    
  2. Elimina Cloud TPU.

    $ gcloud compute tpus tpu-vm delete tpu-name \
      --zone europe-west4-a
    
  3. Verifica che le risorse siano state eliminate eseguendo il comando seguente. Assicurati che il tuo TPU non sia più presente nell'elenco. L'eliminazione può richiedere qualche minuto.

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