Esegui il codice TensorFlow nelle sezioni di pod di TPU

Questo documento mostra come eseguire un calcolo utilizzando TensorFlow su un pod di TPU. Esegui i seguenti passaggi:

  1. Creare un segmento di pod di TPU con il software TensorFlow
  2. Connettiti alla VM TPU tramite SSH
  3. Crea ed esegui uno script di esempio

La VM TPU si basa su account di servizio per le autorizzazioni per chiamare l'API Cloud TPU. Per impostazione predefinita, la VM TPU utilizzerà l'account di servizio predefinito di Compute Engine, che include tutte le autorizzazioni Cloud TPU necessarie. Se utilizzi il tuo account di servizio, devi aggiungere il ruolo Visualizzatore TPU al tuo account di servizio. Per saperne di più sui ruoli di Google Cloud, consulta Informazioni sui ruoli. Puoi specificare il tuo account di servizio utilizzando il flag --service-account quando crei la VM TPU.

Configura l'ambiente

  1. In Cloud Shell, esegui il seguente comando per assicurarti di eseguire la versione corrente di gcloud:

    $ gcloud components update

    Se devi installare gcloud, utilizza il seguente comando:

    $ sudo apt install -y google-cloud-sdk

  2. Crea alcune variabili di ambiente:

    $ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=europe-west4-a
$ export RUNTIME_VERSION=tpu-vm-tf-2.18.0-pod-pjrt
$ export ACCELERATOR_TYPE=v3-32

Creare un segmento di pod di TPU v3-32 con il runtime TensorFlow

$ gcloud compute tpus tpu-vm create ${TPU_NAME}} \
  --zone=${ZONE} \
  --accelerator-type=${ACCELERATOR_TYPE} \
  --version=${RUNTIME_VERSION}

Descrizioni dei flag dei comandi

zone
La zona in cui prevedi di creare la Cloud TPU.
accelerator-type
Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni TPU.
version
La versione software di Cloud TPU.

Connettiti alla VM Cloud TPU tramite SSH

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --zone=${ZONE}

Crea ed esegui uno script di esempio

  1. Imposta le seguenti variabili di ambiente.

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

  2. Crea un file denominato tpu-test.py nella directory corrente e copia e incolla al suo interno il seguente script.

    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. Esegui questo script con il seguente comando:

    (vm)$ python3 tpu-test.py

    Questo script esegue un calcolo su ogni TensorCore di uno slice di pod di TPU. L'output sarà simile al seguente:

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

Esegui la pulizia

Al termine dell'utilizzo della VM TPU, segui questi passaggi per ripulire le risorse.

  1. Disconnettiti da Compute Engine:

    (vm)$ exit

  2. Elimina la Cloud TPU.

    $ gcloud compute tpus tpu-vm delete ${TPU_NAME} \
  --zone=${ZONE}

  3. Verifica che le risorse siano state eliminate eseguendo il seguente comando. Assicurati che la TPU non sia più inclusa nell'elenco. L'eliminazione può richiedere qualche minuto.

    $ gcloud compute tpus tpu-vm list \
  --zone=${ZONE}