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

Este documento mostra como realizar um cálculo usando o TensorFlow em um pod de TPU. Você vai realizar as seguintes etapas:

  1. Criar uma fração do pod de TPU com o software TensorFlow
  2. Conectar-se à VM de TPU usando SSH
  3. Criar e executar um script de exemplo

A VM do TPU depende de uma conta de serviço para ter permissão para chamar a API Cloud TPU. Por padrão, a VM do TPU usa a conta de serviço padrão do Compute Engine, que inclui todas as permissões necessárias do Cloud TPU. Se você usar sua própria conta de serviço, adicione o papel Visualizador de TPU à conta de serviço. Para mais informações sobre papéis do Google Cloud, consulte Noções básicas sobre papéis. É possível especificar sua própria conta de serviço usando a flag --service-account ao criar a VM de TPU.

Configurar o ambiente

  1. No Cloud Shell, execute o seguinte comando para garantir que você está executando a versão atual do gcloud:

    $ gcloud components update

    Se você precisar instalar o gcloud, use o seguinte comando:

    $ sudo apt install -y google-cloud-sdk

  2. Crie algumas variáveis de 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

Criar uma fração do pod de TPU v3-32 com o ambiente de execução do TensorFlow

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

Descrições de sinalizações de comando

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 software do Cloud TPU.

Conectar-se à VM do Cloud TPU usando SSH

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

Criar e executar um script de exemplo

  1. Defina as seguintes variáveis de ambiente.

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

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

  3. Execute esse script com o seguinte comando:

    (vm)$ python3 tpu-test.py

    Esse script realiza um cálculo em cada TensorCore de uma fatia do pod de 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=${ZONE}

  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=${ZONE}