Guia de início rápido: executar um cálculo em uma VM da Cloud TPU usando o TensorFlow

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

Neste guia de início rápido, você verá como criar um Cloud TPU, instalar o TensorFlow e executar um cálculo simples em um Cloud TPU. Para ver um tutorial mais detalhado sobre 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, crie uma conta do Google Cloud Platform, instale a Google Cloud CLI e configure o comando gcloud. Para mais informações, consulte Configurar uma conta e um projeto do Cloud TPU.

Crie uma VM ou nó do Cloud TPU com gcloud

Inicie uma Cloud TPU do Compute Engine usando o comando gcloud. O comando usado depende se você usa uma VM ou um nó de TPU. Para mais informações sobre a arquitetura de duas VMs, consulte Arquitetura do sistema. Para mais informações sobre o comando gcloud, consulte a referência de gcloud.

VM de TPU

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

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.

Nó da TPU

$ gcloud compute tpus execution-groups create \
--name=tpu-name \
--zone=europe-west4-a \
--disk-size=300 \
--machine-type=n1-standard-16 \
--tf-version=2.11.0 \

Descrições de sinalizações de comando

project
ID do projeto do GCP
name
O nome do Cloud TPU a ser criado.
zone
A zona em que você planeja criar a Cloud TPU.
disk-size
O tamanho do disco rígido em GB da VM criada pelo comando gcloud.
machine-type
O tipo de máquina da VM do Compute Engine a ser criada.
tf-version
A versão do TensorFlow gcloud é instalada na VM. Consulte
accelerator-type
O tipo do Cloud TPU a ser criado.

Conecte-se à VM do Cloud TPU

Ao usar VMs de TPU, é necessário usar o SSH explicitamente na VM da TPU. Ao usar os nós da TPU, o SSH deve ser conectado automaticamente à sua VM do Compute Engine. Se a conexão não for automática, use o comando a seguir.

VM de TPU

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

Nó da TPU

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

Execute um exemplo simples usando o TensorFlow

VM de TPU

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

(vm)$ export TPU_NAME=local

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)

Nó da TPU

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

tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='your-tpu-name')  # TPU detection
print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

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

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

Limpar

Para evitar cobranças na sua 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.

    VM de TPU

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

    Nó de TPU

    $ gcloud compute tpus execution-groups 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.

    VM de TPU

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

    Nó da TPU

    $ gcloud compute tpus execution-groups list --zone=europe-west4-a

A seguir

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