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

Neste documento, mostramos como realizar um cálculo simples usando o TensorFlow em um pod de TPU. Você executará as seguintes etapas:

Criar uma fração do Pod de TPU com o software do TensorFlow
Conectar-se à VM da TPU usando SSH
Criar e executar um script simples

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

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=europe-west4-a \
  --accelerator-type=v3-32 \
  --version=tpu-vm-tf-2.14.0-pod

Descrições de sinalizações de comando

zone: A zona em que você planeja criar o 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.

Conectar-se à VM do Cloud TPU usando SSH

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

Criar e executar um script de cálculo simples

Defina as seguintes variáveis de ambiente.

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

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)

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

Limpeza

Quando terminar de usar a VM de TPU, siga estas etapas para limpar os recursos.

Desconecte-se do Compute Engine:
```
(vm)$ exit
```

Exclua o Cloud TPU.

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

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 europe-west4-a
```