Execute um cálculo numa VM de TPU na nuvem com o PyTorch

Crie uma Cloud TPU com gcloud

1. Defina algumas variáveis de ambiente para facilitar a utilização 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 das variáveis de ambiente

   Variável	Descrição
   PROJECT_ID	O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
   TPU_NAME	O nome da TPU.
   ZONE	A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
   ACCELERATOR_TYPE	O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
   RUNTIME_VERSION	A versão do software do Cloud TPU.

2. 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

Estabeleça ligação à sua VM da Cloud TPU

Estabeleça ligação à VM de TPU através de SSH com o seguinte comando:

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

Se não conseguir estabelecer ligação a uma VM de TPU através de SSH, pode dever-se ao facto de a VM de TPU não ter um endereço IP externo. Para aceder a uma VM da TPU sem um endereço IP externo, siga as instruções em Estabeleça ligação a uma VM da TPU sem um endereço IP público.

Instale o PyTorch/XLA na sua 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] -f https://storage.googleapis.com/libtpu-releases/index.html

Verifique se o PyTorch consegue aceder às TPUs

Use o seguinte comando para verificar se o PyTorch consegue aceder às suas TPUs:

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

O resultado do comando deve ter o seguinte aspeto:

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

Fazer um cálculo básico

1. Crie um ficheiro denominado tpu-test.py no diretório atual e copie e cole o seguinte script no mesmo:

   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)
   

2. Execute o script:

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

   O resultado 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')