Esta página foi traduzida pela API Cloud Translation.

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

Neste guia de início rápido, mostramos como criar um Cloud TPU, instalar o PyTorch e executar um cálculo simples em um Cloud TPU. Para conferir um tutorial mais detalhado mostrando como treinar um modelo em um Cloud TPU, consulte um dos Tutoriais do Cloud TPU PyTorch.

Antes de começar

Antes de seguir este guia de início rápido, é preciso 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.

Criar um Cloud TPU com `gcloud`

Para criar uma VM da TPU no projeto, na rede e na zona de usuário padrão, faça o seguinte:

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central1-a \
   --accelerator-type=v3-8 \
   --version=tpu-ubuntu2204-base

Descrições de sinalizações de comando

zone: A zona em que você planeja criar a Cloud TPU.
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 com cada versão de TPU, consulte Versões de TPU.
version: A versão do software do Cloud TPU.

Ao criar a TPU, você pode transmitir as flags --network e --subnetwork adicionais se quiser especificar a rede e a sub-rede padrão. Se você não quiser usar a rede padrão, transmita a flag --network. A flag --subnetwork é opcional e pode ser usada para especificar uma sub-rede padrão para qualquer rede que você esteja usando (padrão ou especificada pelo usuário). Consulte a página de referência da API gcloud para saber mais sobre essas flags.

Conecte-se à VM do Cloud TPU

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-a

Instalar o PyTorch/XLA na VM da TPU

   (vm)$ pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html

Definir a configuração do ambiente de execução da TPU

Verifique se o ambiente de execução do PyTorch/XLA usa a TPU.

   (vm) $ export PJRT_DEVICE=TPU

Verificar se o PyTorch pode acessar TPUs

Crie um arquivo chamado tpu-count.py no diretório atual e copie e cole o script a seguir nele.

import torch
import torch_xla
import torch_xla.core.xla_model as xm
print(f'PyTorch can access {len(torch_xla.devices())} TPU cores')

Execute o script:
```
(vm)$ python3 tpu-count.py
```
A saída do script mostra o resultado do cálculo:
```
PyTorch can access 8 TPU cores
```

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)$ 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 \
  --zone=us-central1-a

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