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 simple en una Cloud TPU. Para obtener un instructivo más detallado que 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, 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 o nodo de Cloud TPU con gcloud
Inicia una Cloud TPU de Compute Engine 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 compute tpus tpu-vm create tpu-name \
--zone=europe-west4-a \
--accelerator-type=v3-8 \
--version=tpu-vm-tf-2.11.0
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.
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.11.0 \
Descripciones de las marcas de comandos
project
- El ID de tu proyecto de GCP
name
- El nombre de la Cloud TPU que se creará.
zone
- La zona en la que planeas crear tu Cloud TPU.
disk-size
- El tamaño del disco duro en GB de la VM que creó 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. Consulta 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 tener una conexión SSH automática a tu VM de Compute Engine. Si no te conectas de forma automática, usa el siguiente comando.
VM de TPU
$ gcloud 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
Una vez que te conectes a la VM de TPU, configura la siguiente variable de entorno.
(vm)$ export TPU_NAME=local
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)
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 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) }
Realiza una limpieza
Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos que se usaron en esta página, sigue estos pasos.
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.
VM de TPU
$ gcloud 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
Ejecuta
gcloud compute tpus tpu-vm list
para verificar que los recursos se hayan borrado. La eliminación puede tardar varios minutos.VM de TPU
$ gcloud 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: