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 y ejecutar un cálculo en una Cloud TPU con TensorFlow. Para obtener un instructivo más detallado que te muestre cómo entrenar un modelo en 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 ver más consulta Configura una cuenta y un proyecto de Cloud TPU.

Crea una VM de Cloud TPU con gcloud

Crea una Cloud TPU con el comando gcloud.

  $ gcloud compute tpus tpu-vm create tpu-name 

    --zone=europe-west4-a 

    --accelerator-type=v3-8 

    --version=tpu-vm-tf-2.17.0-pjrt 

    --project=your-gcp-project-name

Descripciones de las marcas de comandos

tpu-name
Es el nombre de la Cloud TPU que se creará.
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 del entorno de ejecución de TPU. La versión que se usa en esta guía de inicio rápido tiene TensorFlow preinstalado.
project
Es el nombre del proyecto de Google Cloud CLI en el que crearás tu Cloud TPU.

Para obtener más información sobre el comando gcloud, consulta la referencia de gcloud.

Conéctate a tu VM de Cloud TPU

Conéctate a tu VM de TPU con SSH:

  $ gcloud compute tpus tpu-vm ssh tpu-name 

      --zone europe-west4-a 

      --project=your-gcp-project-name

Verifica que TensorFlow pueda acceder a las TPU

  1. Crea un archivo llamado tpu-count.py en el directorio actual y cópialo pega la siguiente secuencia de comandos.

    import tensorflow as tf

print(f'TensorFlow can access {len(tf.config.list_logical_devices('TPU'))} TPU cores')

  2. Ejecuta la secuencia de comandos:

    (vm)$ python3 tpu-count.py

    El resultado de la secuencia de comandos muestra la cantidad de núcleos de TPU disponibles para la VM de TPU:

    TensorFlow can access 8 TPU cores

Ejecuta un procesamiento básico 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 estableces el parámetro --version en una versión que termina en -pjrt, establece 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 usaste en esta página.

  1. 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.

  2. Borra tu Cloud TPU.

      $ gcloud compute tpus tpu-vm delete tpu-name 
    
    --zone=europe-west4-a

  3. Ejecuta 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: