Esegui un calcolo su una VM Cloud TPU utilizzando PyTorch

Questo documento fornisce una breve introduzione all'utilizzo di PyTorch e Cloud TPU.

Prima di iniziare

Prima di eseguire i comandi in questo documento, devi creare un account Google Cloud , installare Google Cloud CLI e configurare il comando gcloud. Per maggiori informazioni, consulta Configurare l'ambiente Cloud TPU.

Crea una Cloud TPU utilizzando gcloud

  1. Definisci alcune variabili di ambiente per semplificare l'utilizzo dei comandi.

    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

    Descrizioni delle variabili di ambiente

    PROJECT_ID
    Il tuo Google Cloud ID progetto.
    ACCELERATOR_TYPE
    Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni di TPU.
    ZONE
    La zona in cui prevedi di creare la tua Cloud TPU.
    RUNTIME_VERSION
    La versione del runtime Cloud TPU.
    TPU_NAME
    Il nome assegnato dall'utente alla tua Cloud TPU.

  2. Crea la VM TPU eseguendo il seguente comando:

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

Connettiti alla VM Cloud TPU

Connettiti alla VM TPU tramite SSH utilizzando il seguente comando:

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

Installa PyTorch/XLA sulla tua VM 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

Verificare che PyTorch possa accedere alle TPU

Utilizza il seguente comando per verificare che PyTorch possa accedere alle TPU.

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

L'output del comando dovrebbe essere simile al seguente:

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

Eseguire un calcolo di base

  1. Crea un file denominato tpu-test.py nella directory corrente e copia e incolla al suo interno il seguente script.

    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. Esegui lo script:

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

    L'output dello script mostra il risultato del calcolo:

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

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questa pagina, segui questi passaggi.

  1. Se non l'hai già fatto, scollega l'istanza Compute Engine:

    (vm)$ exit

    Il tuo prompt dovrebbe ora essere username@projectname, a indicare che ti trovi in Cloud Shell.

  2. Elimina la Cloud TPU.

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

L'output di questo comando dovrebbe confermare che la TPU è stata eliminata.

Passaggi successivi

Scopri di più sulle VM Cloud TPU: