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 che mostra come per addestrare un modello su una Cloud TPU, guarda 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 per 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.17.0-pjrt
    --project=your-gcp-project-name

Descrizioni flag di comando

tpu-name
Il nome della Cloud TPU da creare.
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 del runtime TPU. TensorFlow è preinstallato nella versione utilizzata in questa guida rapida.
project
Il nome del progetto Google Cloud CLI in cui stai creando la 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

Esegui un esempio utilizzando TensorFlow

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

  (vm)$ export TPU_NAME=local

Durante la creazione della TPU, se imposti il parametro --version su una versione che termina con -pjrt, imposta le seguenti variabili di ambiente per abilitare 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, quindi copia e incolla 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 per le risorse utilizzate in questa pagina, segui questi passaggi.

  1. Disconnettiti dall'istanza Compute Engine, se non lo hai già fatto Fatto:

    (vm)$ exit

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

  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 gcloud compute tpus tpu-vm list. La l'eliminazione potrebbe richiedere diversi minuti.

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

Passaggi successivi

Per ulteriori informazioni su Cloud TPU, vedi: