Esegui il codice TensorFlow su sezioni di pod di TPU

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

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

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 utilizzerà account di servizio Compute Engine predefinito che include tutte le autorizzazioni di Cloud TPU necessarie. Se utilizzi il tuo Account di servizio, devi aggiungere il visualizzatore TPU. al tuo account di servizio. Per saperne di più sui ruoli di Google Cloud, consulta Informazioni sui ruoli. Puoi specificare il tuo account di servizio utilizzando il flag --service-account quando durante la creazione della VM TPU.

Crea una sezione di 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.17.0-pod-pjrt

Descrizioni flag di comando

zone
La zona in cui prevedi di creare la tua Cloud TPU.
accelerator-type
Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, vedi Versioni TPU.
version
La versione software di Cloud TPU.

Connettiti alla 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 seguenti variabili di ambiente.

    (vm)$ export TPU_NAME=tpu-name
    (vm)$ export TPU_LOAD_LIBRARY=0
    
  2. Crea un file denominato tpu-test.py nella directory corrente, quindi copia e incolla 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 lo script con il comando seguente:

    (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

Al termine delle operazioni con la VM TPU, segui questi passaggi per pulire le risorse.

  1. Disconnettiti da Compute Engine:

    (vm)$ exit
    
  2. Elimina il tuo Cloud TPU.

    $ gcloud compute tpus tpu-vm delete tpu-name \
      --zone europe-west4-a
    
  3. Verifica che le risorse siano state eliminate eseguendo questo comando. Marca assicurati che la tua TPU non sia più elencata. L'eliminazione può richiedere qualche minuto.

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