Guia de início rápido do TensorFlow da Cloud TPU

Neste guia de início rápido, mostramos como criar um Cloud TPU, instalar o TensorFlow e executar um cálculo simples em um Cloud TPU. Para ver um tutorial mais detalhado mostrando como treinar um modelo em uma Cloud TPU, consulte um dos Tutoriais do Cloud TPU para TensorFlow (em inglês).

Antes de começar

Antes de seguir este guia de início rápido, é preciso criar uma conta do Google Cloud Platform, instalar o SDK do Google Cloud Platform. 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

Use o comando gcloud para iniciar uma VM do Compute Engine e a Cloud TPU. O comando usado depende de uma VM ou um nó de TPU. Para ver 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 alpha compute tpus tpu-vm create tpu-name \
--zone=europe-west4-a \
--accelerator-type=v3-8 \
--version=v2-alpha

Descrições de sinalizações de comando

zone
A zona em que você planeja criar a Cloud TPU.
accelerator-type
O tipo da Cloud TPU a ser criada.
version
A versão do ambiente de execução 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.6.0 \
--accelerator-type=v3-8

Descrições de sinalizações de comando

project
O ID do projeto do GCP
name
O nome da Cloud TPU a ser criada.
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 gcloud do Tensorflow é instalada na VM.
accelerator-type
O tipo da Cloud TPU a ser criada.

Conecte-se à VM da Cloud TPU

Ao usar VMs de TPU, é necessário usar o SSH explicitamente na VM da TPU. Ao usar nós da TPU, o SSH no Compute Engine será automaticamente aplicado à sua VM no Compute Engine. Se a conexão não for estabelecida automaticamente, use o comando a seguir.

VM de TPU

$ gcloud alpha 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

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(tpu='local')
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 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

Esse script realiza um cálculo simples em cada núcleo de uma TPU. A saída exibe a saída de cada núcleo de TPU.

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 alpha compute tpus tpu-vm delete tpu-name \
    --zone=europe-west4-a
    

    Nó da TPU

    $ gcloud compute tpus execution-groups delete tpu-name \
    --zone=europe-west4-a
    
  3. Execute gcloud alpha 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 alpha 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: