Mantieni tutto organizzato con le raccolte Salva e classifica i contenuti in base alle tue preferenze.

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.

    $ 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