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.

Ruoli obbligatori

Per ottenere le autorizzazioni necessarie per creare una TPU e connetterti a quest'ultima tramite SSH, chiedi all'amministratore di concederti i seguenti ruoli IAM sul progetto:

Per saperne di più sulla concessione dei ruoli, consulta Gestisci l'accesso a progetti, cartelle e organizzazioni.

Potresti anche riuscire a ottenere le autorizzazioni richieste tramite i ruoli personalizzati o altri ruoli predefiniti.

Crea una Cloud TPU utilizzando gcloud

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

    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

    Descrizioni delle variabili di ambiente

    Variabile Descrizione
    PROJECT_ID L'ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
    TPU_NAME Il nome della TPU.
    ZONE La zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
    ACCELERATOR_TYPE Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per saperne di più sui tipi di acceleratore supportati per ogni versione di TPU, consulta la sezione Versioni di TPU.
    RUNTIME_VERSION La versione software di 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

Se non riesci a connetterti a una VM TPU tramite SSH, il problema potrebbe essere che la VM TPU non ha un indirizzo IP esterno. Per accedere a una VM TPU senza un indirizzo IP esterno, segui le istruzioni riportate in Connettiti a una VM TPU senza un indirizzo IP pubblico.

Installa PyTorch/XLA sulla 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] -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', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

Eseguire un calcolo di base

  1. Crea un file denominato tpu-test.py nella directory attuale e copia e incolla 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. Disconnettiti dall'istanza Cloud TPU, se non l'hai già fatto:

    (vm)$ exit

    Il 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

  3. Verifica che le risorse siano state eliminate eseguendo il seguente comando. Assicurati che la TPU non sia più elencata. L'eliminazione può richiedere qualche minuto.

    $ gcloud compute tpus tpu-vm list \
    --zone=$ZONE

Passaggi successivi

Scopri di più sulle VM Cloud TPU: