Ejecuta el código de TensorFlow en porciones de pod de TPU
Configura un pod de VM de TPU que ejecute TensorFlow y ejecute un cálculo
Usa el siguiente procedimiento para configurar un pod que ejecute TensorFlow y ejecute un cálculo.
$ gcloud compute tpus tpu-vm create tpu-name \ --zone=europe-west4-a \ --accelerator-type=v3-32 \ --version=tpu-vm-tf-2.11.0-pod
Descripciones de las marcas de comandos
zone
- La zona en la que planeas crear tu Cloud TPU.
accelerator-type
- El tipo de Cloud TPU que se creará.
version
- La versión del software de Cloud TPU. La versión predeterminada es la versión más reciente del software de TensorFlow.
Conéctate a tu VM de Cloud TPU
Establece una conexión SSH a una VM de TPU en tu pod de Cloud TPU:
$ gcloud compute tpus tpu-vm ssh tpu-name \ --zone europe-west4-a
Configura las siguientes variables de entorno:
(vm)$ export TPU_NAME=tpu-name (vm)$ export TPU_LOAD_LIBRARY=0
Crea un archivo llamado
tpu-test.py
en el directorio actual y copia y pega la siguiente secuencia de comandos en él.
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)
Ejecuta esta secuencia de comandos con el siguiente comando:
(vm)$ python3 tpu-test.py
Esta secuencia de comandos realiza un cálculo simple en cada TensorCore de una TPU. El resultado se verá similar al siguiente código:
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) }
Realiza una limpieza
Cuando termines de usar tu VM de TPU, sigue estos pasos para limpiar tus recursos.
Desconéctate de Compute Engine:
(vm)$ exit
Borra tu Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone europe-west4-a
Verifica que se hayan borrado los recursos mediante la ejecución del siguiente comando. Asegúrate de que tu TPU ya no aparezca en la lista. La eliminación puede tardar varios minutos.
$ gcloud compute tpus tpu-vm list \ --zone europe-west4-a