Esegui il codice TensorFlow su sezioni di pod di TPU
Questo documento mostra come eseguire un calcolo semplice utilizzando TensorFlow su di un pod di TPU. Devi seguire questi passaggi:
- Creazione di una sezione di pod di TPU con il software TensorFlow
- Connettiti alla VM TPU tramite SSH
- Creare ed eseguire un semplice script
La VM TPU si basa su un account di servizio
per le autorizzazioni
per chiamare l'API Cloud TPU. Per impostazione predefinita, la VM TPU utilizzerà
account di servizio Compute Engine predefinito
che include tutte le autorizzazioni
di Cloud TPU necessarie. Se utilizzi il tuo
Account di servizio, devi aggiungere il visualizzatore TPU.
al tuo 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
quando
durante la creazione della VM TPU.
Crea una sezione di pod di TPU v3-32 con runtime TensorFlow
$ gcloud compute tpus tpu-vm create tpu-name \ --zone=europe-west4-a \ --accelerator-type=v3-32 \ --version=tpu-vm-tf-2.17.0-pod-pjrt
Descrizioni flag di comando
zone
- La zona in cui prevedi di creare la tua Cloud TPU.
accelerator-type
- Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, vedi 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
Crea ed esegui un semplice script di calcolo
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 corrente, quindi copia e incolla 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 lo 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 con la VM TPU, segui questi passaggi per pulire le risorse.
Disconnettiti da Compute Engine:
(vm)$ exit
Elimina il tuo Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone europe-west4-a
Verifica che le risorse siano state eliminate eseguendo questo comando. Marca 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