Esegui il codice TensorFlow sulle sezioni di pod TPU
Configura un pod della VM TPU che esegue TensorFlow ed esegui un calcolo
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.
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
Imposta le seguenti variabili di ambiente.
(vm)$ export TPU_NAME=tpu-name (vm)$ export TPU_LOAD_LIBRARY=0
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.
Disconnetti da Compute Engine:
(vm)$ exit
Elimina la tua Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone europe-west4-a
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