Esegui il codice TensorFlow sulle sezioni di pod di TPU

Questo documento mostra come eseguire un semplice calcolo utilizzando TensorFlow su un pod di TPU. Dovrai procedere nel seguente modo:

  1. Crea una sezione di pod di TPU con il software TensorFlow
  2. Connettiti alla VM TPU tramite SSH
  3. Creare ed eseguire uno script semplice

La VM TPU si basa su un account di servizio per le autorizzazioni di chiamata all'API Cloud TPU. Per impostazione predefinita, la VM TPU utilizzerà l'account di servizio predefinito di Compute Engine, che include tutte le autorizzazioni Cloud TPU necessarie. Se utilizzi un tuo account di servizio, devi aggiungere il ruolo Visualizzatore TPU all'account di servizio. Per saperne di più sui ruoli di Google Cloud, consulta Informazioni sui ruoli. Puoi specificare il tuo account di servizio utilizzando il flag --service-account durante la creazione della VM TPU.

Crea una sezione di pod di TPU v3-32 con il runtime TensorFlow

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=europe-west4-a \
  --accelerator-type=v3-32 \
  --version=tpu-vm-tf-2.16.1-pod-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.

Connettiti alla VM Cloud TPU tramite SSH

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

Creare ed eseguire un semplice script di calcolo

  1. Imposta le seguenti variabili di ambiente.

    (vm)$ export TPU_NAME=tpu-name
(vm)$ export TPU_LOAD_LIBRARY=0

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

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
print('Running on TPU ', cluster_resolver.cluster_spec().as_dict()['worker'])

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

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

  3. 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),
8: tf.Tensor(2.0, shape=(), dtype=float32),
9: tf.Tensor(2.0, shape=(), dtype=float32),
10: tf.Tensor(2.0, shape=(), dtype=float32),
11: tf.Tensor(2.0, shape=(), dtype=float32),
12: tf.Tensor(2.0, shape=(), dtype=float32),
13: tf.Tensor(2.0, shape=(), dtype=float32),
14: tf.Tensor(2.0, shape=(), dtype=float32),
15: tf.Tensor(2.0, shape=(), dtype=float32),
16: tf.Tensor(2.0, shape=(), dtype=float32),
17: tf.Tensor(2.0, shape=(), dtype=float32),
18: tf.Tensor(2.0, shape=(), dtype=float32),
19: tf.Tensor(2.0, shape=(), dtype=float32),
20: tf.Tensor(2.0, shape=(), dtype=float32),
21: tf.Tensor(2.0, shape=(), dtype=float32),
22: tf.Tensor(2.0, shape=(), dtype=float32),
23: tf.Tensor(2.0, shape=(), dtype=float32),
24: tf.Tensor(2.0, shape=(), dtype=float32),
25: tf.Tensor(2.0, shape=(), dtype=float32),
26: tf.Tensor(2.0, shape=(), dtype=float32),
27: tf.Tensor(2.0, shape=(), dtype=float32),
28: tf.Tensor(2.0, shape=(), dtype=float32),
29: tf.Tensor(2.0, shape=(), dtype=float32),
30: tf.Tensor(2.0, shape=(), dtype=float32),
31: tf.Tensor(2.0, shape=(), dtype=float32)
}

Esegui la pulizia

Al termine delle operazioni della VM TPU, segui questi passaggi per la pulizia delle risorse.

  1. Disconnettiti da Compute Engine:

    (vm)$ exit

  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 questo comando. Assicurati che la tua TPU non sia più elencata. L'eliminazione può richiedere qualche minuto.

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