Esegui un calcolo su una VM Cloud TPU utilizzando TensorFlow

Questa guida rapida mostra come creare una Cloud TPU ed eseguire un calcolo su una Cloud TPU utilizzando TensorFlow. Per un tutorial più approfondito su come addestrare un modello su una Cloud TPU, consulta uno dei tutorial su Cloud TPU.

Prima di iniziare

Prima di seguire questa guida rapida, devi creare un account Google Cloud, installare Google Cloud CLI e configurare il comando gcloud. Per maggiori informazioni, consulta Configurare un account e un progetto Cloud TPU.

Crea una VM Cloud TPU con gcloud

Crea una Cloud TPU utilizzando il comando gcloud.

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=europe-west4-a \
  --accelerator-type=v3-8 \
  --version=tpu-vm-tf-2.18.0-pjrt \
  --project=your-gcp-project-name

Descrizioni dei flag dei comandi

tpu-name
Il nome della Cloud TPU da creare.
zone
La zona in cui prevedi di creare la 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, consulta Versioni TPU.
version
La versione del runtime TPU. La versione utilizzata in questa guida rapida ha TensorFlow preinstallato.
project
Il nome del progetto Google Cloud CLI in cui stai creando Cloud TPU.

Per ulteriori informazioni sul comando gcloud, consulta la pagina di riferimento a gcloud.

Connettiti alla VM Cloud TPU

Connettiti alla VM TPU tramite SSH:

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --zone europe-west4-a \
  --project=your-gcp-project-name

Verificare che TensorFlow possa accedere alle TPU

  1. Crea un file denominato tpu-count.py nella directory corrente e copia e incolla al suo interno il seguente script.

    import tensorflow as tf

print(f"TensorFlow can access {len(tf.config.list_logical_devices('TPU'))} TPU cores")

  2. Esegui lo script:

    (vm)$ python3 tpu-count.py

    L'output dello script mostra il numero di core TPU disponibili per la VM TPU:

    TensorFlow can access 8 TPU cores

Eseguire un calcolo di base utilizzando TensorFlow

Dopo aver eseguito la connessione alla VM TPU, imposta la seguente variabile di ambiente.

  (vm)$ export TPU_NAME=local

Quando crei la TPU, se imposti il parametro --version su una versione che termina con -pjrt, imposta le seguenti variabili di ambiente per attivare il runtime PJRT: 

  (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

Crea un file denominato tpu-test.py nella directory corrente e copia e incolla al suo interno il seguente script.

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)

Esegui questo script con il seguente comando:

(vm)$ python3 tpu-test.py

Questo script esegue un 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)
}

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questa pagina, segui questi passaggi.

  1. Se non l'hai già fatto, disconnetti dall'istanza Compute Engine:

    (vm)$ exit

    Il tuo prompt dovrebbe ora essere username@projectname, a indicare che ti trovi in Cloud Shell.

  2. Elimina la Cloud TPU.

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

  3. Verifica che le risorse siano state eliminate eseguendo gcloud compute tpus tpu-vm list. L'eliminazione potrebbe richiedere alcuni minuti.

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

Passaggi successivi

Per ulteriori informazioni su Cloud TPU, consulta: