Esegui un calcolo su una VM Cloud TPU utilizzando TensorFlow

Questa guida rapida mostra come creare una Cloud TPU, installare TensorFlow ed eseguire un calcolo su una Cloud TPU. Per un tutorial più approfondito che mostra come addestrare un modello su una Cloud TPU, consulta uno dei tutorial per 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 saperne di più, vedi Configurare un account e un progetto Cloud TPU.

Crea una VM Cloud TPU con gcloud

Avvia una Cloud TPU di Compute Engine utilizzando il comando gcloud. Per ulteriori informazioni sul comando gcloud, consulta il riferimento gcloud.

Quando avvii la VM TPU, puoi utilizzare la versione software TPU predefinita mostrata nel comando oppure puoi fare riferimento alle versioni software Cloud TPU per specificare un'altra versione software TPU supportata.

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=europe-west4-a
    --accelerator-type=v3-8
    --version=tpu-vm-tf-2.16.1-pjrt

Descrizioni flag comando

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 maggiori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni TPU.
version
La versione del software TPU.

Connettiti alla VM Cloud TPU

Devi connetterti esplicitamente alla VM TPU tramite SSH. Se la connessione non viene eseguita automaticamente, utilizza il comando seguente.

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

Esegui un esempio utilizzando TensorFlow

Una volta stabilita 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 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 comando seguente:

(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. Disconnettiti dall'istanza di Compute Engine, se non lo hai già fatto:

    (vm)$ exit

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

  2. Elimina la tua 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 diversi minuti.

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

Passaggi successivi

Per ulteriori informazioni su Cloud TPU, vedi: