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:
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.Elimina la tua Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name
--zone=europe-west4-aVerifica 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: