Ejecuta el código de TensorFlow en porciones de pod de TPU

En este documento, se muestra cómo realizar un cálculo simple mediante TensorFlow en un pod de TPU. Realizarás los siguientes pasos:

  1. Crea una porción de pod de TPU con el software de TensorFlow
  2. Conéctate a la VM de TPU con SSH
  3. Crea y ejecuta una secuencia de comandos simple

La VM de TPU se basa en una cuenta de servicio para obtener 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 propia cuenta de servicio, debes agregar el rol de visualizador 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 crees tu VM de TPU.

Crea 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.16.1-pod-pjrt

Descripciones de las marcas de comandos

zone
La zona en la que planeas crear tu Cloud TPU.
accelerator-type
El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que quieres crear. Si quieres 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 simple

  1. Configura las siguientes variables de entorno:

    (vm)$ export TPU_NAME=tpu-name
(vm)$ export TPU_LOAD_LIBRARY=0

  2. 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)

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

  1. Desconéctate de Compute Engine:

    (vm)$ exit

  2. Borra tu Cloud TPU.

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

  3. Ejecuta el siguiente comando para verificar que se hayan borrado los recursos. Asegúrate 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