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

Este documento apresenta uma breve introdução ao trabalho com o PyTorch e o Cloud TPU.

Antes de começar

Antes de executar os comandos neste documento, crie uma conta do Google Cloud , instale a CLI do Google Cloud e configure o comando gcloud. Para mais informações, consulte Configurar o ambiente do Cloud TPU.

Criar um Cloud TPU usando `gcloud`

Defina algumas variáveis de ambiente para facilitar o uso dos comandos.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east5-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Descrições de variáveis de ambiente

Variável	Descrição
`PROJECT_ID`	O ID do projeto do Google Cloud . Use um projeto atual ou crie um novo.
`TPU_NAME`	O nome da TPU.
`ZONE`	A zona em que a VM de TPU será criada. Para mais informações sobre as zonas disponíveis, consulte Zonas e regiões de 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 aceitos por cada versão de TPU, consulte Versões de TPU.
`RUNTIME_VERSION`	A versão do software do Cloud TPU.

Crie a VM de TPU executando o seguinte comando:

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

Conectar-se à VM do Cloud TPU

Conecte-se à VM de TPU por SSH usando o seguinte comando:

$ gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

Se você não consegue se conectar a uma VM de TPU usando SSH, ela pode não ter um endereço IP externo. Para acessar uma VM de TPU sem um endereço IP externo, siga as instruções em Conectar-se a uma VM de TPU sem um endereço IP público.

Instalar o PyTorch/XLA na VM de TPU

$ (vm) sudo apt-get update
$ (vm) sudo apt-get install libopenblas-dev -y
$ (vm) pip install numpy
$ (vm) pip install torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

Verificar se o PyTorch pode acessar as TPUs

Use o comando abaixo para verificar se o PyTorch pode acessar as TPUs:

$ (vm) PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

A saída do comando será parecida com esta:

['xla:0', 'xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

Fazer um cálculo básico

Crie um arquivo tpu-test.py no diretório atual e copie e cole o script abaixo nele.

import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

Execute o script:

(vm)$ PJRT_DEVICE=TPU python3 tpu-test.py

A saída do script mostra o resultado do cálculo:

tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

Limpeza

Para evitar cobranças na conta do Google Cloud pelos recursos usados nesta página, siga as etapas abaixo.

Desconecte-se da instância do Cloud TPU, 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.

$ gcloud compute tpus tpu-vm delete $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

Execute o comando abaixo 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
```

A seguir

Saiba mais sobre as VMs do Cloud TPU: