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

En este documento, se muestra cómo realizar un cálculo simple con 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 mediante SSH
  3. Crea y ejecuta una secuencia de comandos simple

La VM de TPU se basa en cuentas 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 la función 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 creas la 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.15.0-pod-pjrt

Descripciones de las marcas de comandos

zone
La zona en la que planeas crear la Cloud TPU.
accelerator-type
El tipo de Cloud TPU que se creará.
version
La versión de software de Cloud TPU.

Conéctate a la VM de Cloud TPU mediante 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 los recursos se hayan borrado. 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