Esegui il codice TensorFlow su sezioni di pod di TPU

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

  1. Creazione di una sezione 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 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 questo comando per assicurarti di essere con la versione corrente di gcloud:

    $ gcloud components update

    Se devi installare gcloud, usa 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.17.0-pod-pjrt
$ export ACCELERATOR_TYPE=v3-32

Creare un segmento di pod 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 tua 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, vedi 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, quindi copia e incolla 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 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 delle operazioni con la VM TPU, segui questi passaggi per pulire 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 questo 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}