Executar o código do TensorFlow em frações de Pod de TPU

Configurar um pod de VM de TPU executando o TensorFlow e fazer um cálculo

1. Use o procedimento a seguir para configurar um pod que executa o TensorFlow e execute um cálculo.

   Observação: ao iniciar a VM de TPU, é possível usar a versão de software de TPU mostrada no comando ou consultar versões de software de TPU do Cloud para especificar outra versão de software de TPU compatível.

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

   Descrições de sinalizações de comando

   zone

   A zona em que você planeja criar a Cloud TPU.

   accelerator-type

   O tipo do Cloud TPU a ser criado.

   version

   A versão do software do Cloud TPU . O padrão é a versão mais recente do software TensorFlow.

2. Conectar-se à VM do Cloud TPU

   SSH em uma VM da TPU no pod do Cloud TPU:

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

3. Defina as seguintes variáveis de ambiente.

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

4. Crie um arquivo chamado tpu-test.py no diretório atual e copie e cole o script a seguir nele.

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)

Execute esse script com o seguinte comando:

(vm)$ python3 tpu-test.py

Este script executa uma computação simples em cada TensorCore de uma TPU. O resultado será semelhante ao seguinte:

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

Limpar

Quando terminar de usar a VM de TPU, siga estas etapas para limpar os recursos.

1. Desconecte-se do Compute Engine:

   (vm)$ exit
   

2. Exclua o Cloud TPU.

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

3. Execute o seguinte comando para verificar se os recursos foram excluídos. Verifique se a TPU não está mais listada. A exclusão pode levar vários minutos.

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