Executar um cálculo em uma VM da Cloud TPU usando o TensorFlow

Neste guia de início rápido, mostramos como criar um Cloud TPU e executar um cálculo em um Cloud TPU usando o TensorFlow. Para conferir um tutorial mais detalhado mostrando como treinar um modelo em um Cloud TPU, consulte um dos Tutoriais do Cloud TPU.

Antes de começar

Antes de seguir este guia de início rápido, é preciso criar uma conta do Google Cloud, instalar a CLI do Google Cloud e configurar o comando gcloud. Para mais informações, consulte Configurar uma conta e um projeto do Cloud TPU.

Criar uma VM do Cloud TPU com gcloud

Crie uma Cloud TPU usando o comando gcloud.

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=europe-west4-a \
  --accelerator-type=v3-8 \
  --version=tpu-vm-tf-2.18.0-pjrt \
  --project=your-gcp-project-name

Descrições de sinalizações de comando

tpu-name
O nome da Cloud TPU a ser criada.
zone
A zona em que você planeja criar a Cloud TPU.
accelerator-type
O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.
version
A versão do ambiente de execução da TPU. A versão usada neste guia de início rápido tem o TensorFlow pré-instalado.
project
O nome do projeto da Google Cloud CLI em que você está criando a Cloud TPU.

Para mais informações sobre o comando gcloud, consulte a referência do gcloud.

Conecte-se à VM do Cloud TPU

Conecte-se à VM do TPU usando SSH:

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --zone europe-west4-a \
  --project=your-gcp-project-name

Verificar se o TensorFlow pode acessar TPUs

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

    import tensorflow as tf

print(f"TensorFlow can access {len(tf.config.list_logical_devices('TPU'))} TPU cores")

  2. Execute o script:

    (vm)$ python3 tpu-count.py

    A saída do script mostra o número de núcleos de TPU disponíveis para a VM de TPU:

    TensorFlow can access 8 TPU cores

Executar uma computação básica usando o TensorFlow

Depois de se conectar à VM de TPU, defina a seguinte variável de ambiente.

  (vm)$ export TPU_NAME=local

Ao criar a TPU, se você definir o parâmetro --version como uma versão que termina com -pjrt, defina as seguintes variáveis de ambiente para ativar o ambiente de execução do PJRT: 

  (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
  (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

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

@tf.function
def add_fn(x,y):
  z = x + y
  return z

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
strategy = tf.distribute.TPUStrategy(cluster_resolver)

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

Esse script realiza um cálculo 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)
}

Limpar

Para evitar cobranças na conta do Google Cloud pelos recursos usados nesta página, siga estas etapas.

  1. Desconecte-se da instância do Compute Engine, caso ainda não tenha feito isso:

    (vm)$ exit

    Agora, o prompt precisa ser username@projectname, mostrando que você está no Cloud Shell.

  2. Exclua o Cloud TPU.

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

  3. Execute gcloud compute tpus tpu-vm list para verificar se os recursos foram excluídos. A exclusão pode levar vários minutos.

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

A seguir

Para mais informações sobre o Cloud TPU, consulte: