Executar um cálculo em uma VM do 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 um tutorial mais detalhado mostrando treinar um modelo em um Cloud TPU, veja um dos Tutoriais de PyTorch do Cloud TPU.
Antes de começar
Antes de seguir este guia de início rápido, é preciso criar uma conta do Google Cloud Platform, instalar a CLI do Google Cloud e configurar o comando gcloud
.
Para mais informações, consulte
Configure 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 planeja criar o Cloud 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 compatíveis com cada versão de TPU, consulte Versões de TPU.
version
- O software do Cloud TPU versão.
Ao criar sua TPU, é possível transmitir os outros --network
e
--subnetwork
se
você quer especificar a rede e a sub-rede padrão.
Se não quiser usar a rede padrão, você deve passar o valor-chave
sinalização --network
. A sinalização --subnetwork
é opcional e pode ser usada para
especificar uma sub-rede padrão para a rede que estiver usando (padrão ou
(especificado 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
Instale o PyTorch/XLA na VM da TPU
(vm)$ pip install torch~=2.4.0 torch_xla[tpu]~=2.4.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
Faça um cálculo simples:
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.
A seguir
Leia mais sobre VMs do Cloud TPU: