Esegui il codice TensorFlow sulle sezioni di pod TPU

Configura un pod della VM TPU che esegue TensorFlow ed esegui un calcolo

1. Utilizza la procedura seguente per configurare un pod che esegue TensorFlow ed eseguire un calcolo.

   Nota: quando avvii la VM TPU, puoi utilizzare la versione del software TPU 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-32 \
   --version=tpu-vm-tf-2.11.0-pod
   

   Descrizioni flag comando

   zone

   La zona in cui prevedi di creare la tua Cloud TPU.

   accelerator-type

   Il tipo di Cloud TPU da creare.

   version

   La versione del software Cloud TPU . Il valore predefinito è l'ultima versione del software TensorFlow.

2. Connettiti alla VM Cloud TPU

   Accedi a una VM TPU nel tuo pod Cloud TPU:

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

3. Imposta le seguenti variabili di ambiente.

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

4. Crea un file denominato tpu-test.py nella directory attuale, 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)

Esegui questo script con il seguente comando:

(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 della VM TPU, segui questi passaggi per eseguire la pulizia delle risorse.

1. Disconnetti 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 il comando seguente. 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