Ejecuta un cálculo en una VM de Cloud TPU con TensorFlow
En esta guía de inicio rápido, se muestra cómo crear una Cloud TPU, instalar TensorFlow y ejecutar un cálculo en una Cloud TPU. Para obtener un instructivo más detallado en el que se muestra cómo entrenar un modelo en una Cloud TPU, consulta uno de los instructivos de Cloud TPU.
Antes de comenzar
Antes de seguir esta guía de inicio rápido, debes crear una cuenta de Google Cloud, instalar Google Cloud CLI y configurar el comando gcloud
.
Para obtener más información, consulta Configura una cuenta y un proyecto de Cloud TPU.
Crea una VM de Cloud TPU con gcloud
Inicia una Cloud TPU de Compute Engine con el comando gcloud
. Para obtener más información sobre el comando gcloud
, consulta la referencia de gcloud
.
Cuando inicias la VM de TPU, puedes usar la versión predeterminada de software de TPU que se muestra en el comando o puedes consultar las versiones de software de Cloud TPU para especificar otra versión de software de TPU compatible.
$ gcloud compute tpus tpu-vm create tpu-name
--zone=europe-west4-a
--accelerator-type=v3-8
--version=tpu-vm-tf-2.16.1-pjrt
Descripciones de las marcas de comandos
tpu-name
- Es el nombre de la Cloud TPU que se creará.
zone
- 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. Si deseas obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta las versiones de TPU.
version
- La versión del software de la TPU.
Conéctate a tu VM de Cloud TPU
Debes conectarte de forma explícita a la VM de TPU con SSH. Si no te conectas automáticamente, usa el siguiente comando.
$ gcloud compute tpus tpu-vm ssh tpu-name
--zone europe-west4-a
Ejecuta un ejemplo con TensorFlow
Una vez que te conectes a la VM de TPU, configura la siguiente variable de entorno.
(vm)$ export TPU_NAME=local
Cuando creas tu TPU, si configuras el parámetro --version
en una versión que termina en -pjrt
, configura las siguientes variables de entorno para habilitar el entorno de ejecución de PJRT:
(vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
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__) @tf.function def add_fn(x,y): z = x + y return z cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver() tf.config.experimental_connect_to_cluster(cluster_resolver) tf.tpu.experimental.initialize_tpu_system(cluster_resolver) strategy = tf.distribute.TPUStrategy(cluster_resolver) 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 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) }
Limpia
Sigue estos pasos para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos que se usaron en esta página.
Desconéctate de la instancia de Compute Engine, si aún no lo hiciste:
(vm)$ exit
El mensaje ahora debería mostrar
username@projectname
, que indica que estás en Cloud Shell.Borra tu Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name
--zone=europe-west4-aEjecuta
gcloud compute tpus tpu-vm list
para verificar que los recursos se hayan borrado. La eliminación puede tardar varios minutos.$ gcloud compute tpus tpu-vm list --zone=europe-west4-a
¿Qué sigue?
Para obtener más información sobre Cloud TPU, consulta los siguientes vínculos: