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, instalar o TensorFlow e executar um cálculo simples em uma 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, você precisa criar uma conta do Google Cloud Platform, instalar a Google Cloud CLI e configurar 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 do 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.16.1-pjrt
Descrições de sinalizações de comando
zone
- A zona em que você planeja criar o Cloud TPU.
accelerator-type
- O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, acesse Versões de TPU.
version
- A versão do software do Cloud TPU.
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.12.0 \
Descrições de sinalizações de comando
project
- ID do seu projeto do Google Cloud
name
- O nome do Cloud TPU a ser criado.
zone
- A zona em que você planeja criar o 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
gcloud
do TensorFlow é instalada na VM. Consulte accelerator-type
- O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, acesse Versões de TPU.
Conectar-se à VM do Cloud TPU
Ao usar VMs de TPU, conecte-se explicitamente à VM da TPU usando SSH. Ao usar nós de TPU, você precisa automaticamente usar SSH na VM do Compute Engine. Se você não se conectar automaticamente, 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
Executar um exemplo simples com 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
Ao criar sua 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 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)
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 um cálculo 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.
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.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
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: