Esta página foi traduzida pela API Cloud Translation.

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

Este documento apresenta uma breve introdução sobre como trabalhar com o PyTorch e o Cloud TPU.

Antes de começar

Antes de executar os comandos neste documento, é necessário criar uma conta Google Cloud , instalar a Google Cloud CLI e configurar 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
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export TPU_NAME=your-tpu-name
```
Descrições das variáveis de ambiente

PROJECT_ID

O ID do Google Cloud projeto.

ACCELERATOR_TYPE

O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis para cada versão de TPU, consulte Versões de TPU.

ZONE

A zona em que você planeja criar a Cloud TPU.

RUNTIME_VERSION

A versão do ambiente de execução do Cloud TPU.

TPU_NAME

O nome atribuído pelo usuário ao Cloud TPU.

Crie a VM da 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

Conecte-se à VM do Cloud TPU

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

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

Instalar o PyTorch/XLA na VM da 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]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html

Verificar se o PyTorch pode acessar TPUs

Use o comando a seguir para verificar se o PyTorch pode acessar suas 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 vai ficar assim:

['xla:0', 'xla:1', 'xla:2', 'xla:3']

Fazer um cálculo básico

Crie um arquivo chamado tpu-test.py no diretório atual e copie e cole o script a seguir 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')

Limpar

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

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

A saída desse comando confirmará se a TPU foi excluída.