Guía de inicio rápido de TensorFlow de Cloud TPU

En esta guía de inicio rápido, se muestra cómo crear una Cloud TPU, instalar TensorFlow y ejecutar un cálculo simple en Cloud TPU. Para obtener un instructivo más detallado en el que se 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 Platform y, luego, instalar el SDK de Cloud. y configura el comando gcloud. Para obtener más información, consulta Configura una cuenta y un proyecto de Cloud TPU.

Crea una VM o nodo de Cloud TPU con gcloud

Inicia una VM de Compute Engine y Cloud TPU con el comando gcloud. El comando que uses depende de si usas una VM de TPU o un nodo TPU. Para obtener más información sobre la arquitectura de las dos VM, consulta Arquitectura del sistema. Para obtener más información sobre el comando gcloud, consulta la referencia de gcloud.

VM de TPU

$ gcloud alpha compute tpus tpu-vm create tpu-name \
--zone=europe-west4-a \
--accelerator-type=v3-8 \
--version=v2-alpha

Descripciones de las marcas de comandos

zone
Es la zona en la que deseas crear la Cloud TPU.
accelerator-type
El tipo de Cloud TPU que se creará.
version
La versión del entorno de ejecución de Cloud TPU.

Nodo TPU

$ gcloud compute tpus execution-groups create \
--name=tpu-name \
--zone=europe-west4-a \
--disk-size=300 \
--machine-type=n1-standard-16 \
--tf-version=2.6.0 \
--accelerator-type=v3-8

Descripciones de las marcas de comandos

project
: El ID de tu proyecto de GCP
name
El nombre de la Cloud TPU para crear.
zone
Es la zona en la que deseas crear la Cloud TPU.
disk-size
El tamaño del disco duro en GB de la VM creada con el comando gcloud.
machine-type
El tipo de máquina de la VM de Compute Engine que se creará.
tf-version
La versión de Tensorflow gcloud se instala en la VM.
accelerator-type
El tipo de Cloud TPU que se creará.

Conéctate a tu VM de Cloud TPU

Cuando uses las VM de TPU, debes establecer una conexión SSH explícita a tu VM de TPU. Cuando usas nodos TPU, debes conectarte mediante SSH a tu VM de Compute Engine de forma automática. Si no estás conectado de forma automática, usa el siguiente comando.

VM de TPU

$ gcloud alpha compute tpus tpu-vm ssh tpu-name \
  --zone europe-west4-a

Nodo TPU

$ gcloud compute ssh tpu-name \
    --zone=europe-west4-a

Ejecuta un ejemplo simple con TensorFlow

VM de TPU

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(tpu='local')
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)

Nodo TPU

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__)

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='your-tpu-name')  # TPU detection
print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
strategy = tf.distribute.experimental.TPUStrategy(tpu)

@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 núcleo de una TPU. Se muestra el resultado de cada núcleo de TPU.

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.

    VM de TPU

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

    Nodo TPU

    $ gcloud compute tpus execution-groups delete tpu-name \
    --zone=europe-west4-a
    
  3. Ejecuta gcloud alpha compute tpus tpu-vm list para verificar que los recursos se hayan borrado. La eliminación puede tardar varios minutos.

    VM de TPU

    $ gcloud alpha compute tpus tpu-vm list --zone=europe-west4-a
    

    Nodo TPU

    $ gcloud compute tpus execution-groups list --zone=europe-west4-a
    

¿Qué sigue?

Para obtener más información sobre Cloud TPU, consulta los siguientes vínculos: