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

Configura un pod de VM de TPU que ejecute TensorFlow y ejecute un cálculo

1. Usa el siguiente procedimiento para configurar un pod que ejecute TensorFlow y ejecute un cálculo.

   Nota: Cuando inicias la VM de TPU, puedes usar la versión de software de TPU que se muestra en el comando o puedes consultar las versiones de software de Cloud TPU para especificar otra versión de software de TPU compatible.

   $ gcloud compute tpus tpu-vm create tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v3-32 \
   --version=tpu-vm-tf-2.11.0-pod
   

   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.

2. Conéctate a tu VM de Cloud TPU

   Establece una conexión SSH a una VM de TPU en tu pod de Cloud TPU:

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

3. Configura las siguientes variables de entorno:

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

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

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

Realiza una limpieza

Cuando termines de usar tu VM de TPU, sigue estos pasos para limpiar tus 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. Verifica que se hayan borrado los recursos mediante la ejecución del siguiente comando. 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