Esegui un calcolo su una VM Cloud TPU utilizzando TensorFlow

Questa guida rapida mostra come creare una Cloud TPU, installare TensorFlow ed eseguire un semplice 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 per Google Cloud Platform, 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 o un nodo Cloud TPU con gcloud

Avvia una Cloud TPU di Compute Engine utilizzando il comando gcloud. Il comando da utilizzare dipende dall'utilizzo di una VM TPU o di un nodo TPU. Per ulteriori informazioni sull'architettura delle due VM, consulta Architettura di sistema. Per ulteriori informazioni sul comando gcloud, consulta il riferimento gcloud.

VM TPU

$ 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

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 software di Cloud TPU.

Nodo TPU

$ gcloud compute tpus execution-groups create \
--name=tpu-name \
--zone=europe-west4-a \
--disk-size=300 \
--machine-type=n1-standard-16 \
--tf-version=2.12.0 \

Descrizioni flag comando

project
L'ID del tuo progetto Google Cloud
name
Il nome della Cloud TPU da creare.
zone
La zona in cui prevedi di creare la Cloud TPU.
disk-size
Le dimensioni del disco rigido in GB della VM creata dal comando gcloud.
machine-type
Il tipo di macchina della VM di Compute Engine da creare.
tf-version
La versione di TensorFlow gcloud si installa sulla VM. Vedi
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.

Connettiti alla VM Cloud TPU

Se utilizzi le VM TPU, devi connetterti in modo esplicito alla VM TPU tramite SSH. Quando utilizzi i nodi TPU, dovresti connetterti automaticamente tramite SSH alla VM Compute Engine. Se la connessione non viene eseguita automaticamente, utilizza il comando seguente.

VM TPU

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

Nodo TPU

$ gcloud compute ssh tpu-name \
    --zone=europe-west4-a

Esegui un semplice esempio utilizzando TensorFlow

VM TPU

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)

Nodo TPU

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__)

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='your-tpu-name')  # TPU detection
print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.experimental.TPUStrategy(tpu)

@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)

Esegui questo 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)
}

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.

    VM TPU

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

    Nodo TPU

    $ gcloud compute tpus execution-groups 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.

    VM TPU

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

    Nodo TPU

    $ gcloud compute tpus execution-groups list --zone=europe-west4-a

Passaggi successivi

Per ulteriori informazioni su Cloud TPU, vedi: