Ejecuta el código de TensorFlow en porciones de pod de TPU
En este documento, se muestra cómo realizar un cálculo sencillo con TensorFlow en un Pod de TPU. Realizarás los siguientes pasos:
- Crea una porción de pod de TPU con el software de TensorFlow
- Conéctate a la VM de TPU con SSH
- Crea y ejecuta una secuencia de comandos simple
La VM de TPU depende de una cuenta de servicio
los permisos para llamar a la API de Cloud TPU. De forma predeterminada, tu VM de TPU usará la
cuenta de servicio predeterminada de Compute Engine
que incluye todos los permisos necesarios de Cloud TPU. Si usas tu propio plan
cuenta de servicio, debes agregar el Visor de TPU
a tu cuenta de servicio. Para obtener más información sobre los roles de Google Cloud, consulta Comprende los roles.
Puedes especificar tu propia cuenta de servicio con la marca --service-account
cuando
creando tu VM de TPU.
Crear una porción de pod de TPU v3-32 con el entorno de ejecución de 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
Descripciones de las marcas de comandos
zone
- Es la zona en la que deseas crear la Cloud TPU.
accelerator-type
- El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
version
- La versión de software de Cloud TPU.
Conéctate a la VM de Cloud TPU con SSH
$ gcloud compute tpus tpu-vm ssh tpu-name \ --zone europe-west4-a
Crea y ejecuta una secuencia de comandos de cálculo sencilla
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) }
Limpia
Cuando termines de usar la VM de TPU, sigue estos pasos para limpiar los recursos.
Desconéctate de Compute Engine:
(vm)$ exit
Borra tu Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone europe-west4-a
Ejecuta el siguiente comando para verificar que se hayan borrado los recursos. Marca 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