Esegui un calcolo su una VM Cloud TPU utilizzando JAX

Crea una VM 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	Il tuo Google Cloud ID progetto. 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 ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni TPU.
   RUNTIME_VERSION	La versione software di Cloud TPU.

2. Crea la VM TPU eseguendo il seguente comando da Cloud Shell o dal terminale del computer su cui è installato Google Cloud CLI.

   $ 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 utilizzando SSH, il motivo 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 Connettersi a una VM TPU senza un indirizzo IP pubblico.

Installa JAX sulla tua VM Cloud TPU

(vm)$ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Controllo del sistema

Verifica che JAX possa accedere alla TPU ed eseguire operazioni di base:

1. Avvia l'interprete Python 3:

   (vm)$ python3

   >>> import jax

2. Mostra il numero di core TPU disponibili:

   >>> jax.device_count()

Viene visualizzato il numero di core TPU. Il numero di core visualizzato dipende dalla versione della TPU in uso. Per ulteriori informazioni, consulta Versioni TPU.

Eseguire un calcolo

>>> jax.numpy.add(1, 1)

Viene visualizzato il risultato dell'addizione di NumPy:

Output del comando:

Array(2, dtype=int32, weak_type=True)

Uscire dall'interprete Python

>>> exit()

Esecuzione del codice JAX su una VM TPU

Ora puoi eseguire qualsiasi codice JAX. Gli esempi di Flax sono un ottimo punto di partenza per eseguire modelli di ML standard in JAX. Ad esempio, per addestrare una rete convoluzionale MNIST di base:

1. Installa le dipendenze degli esempi Flax:

   (vm)$ pip install --upgrade clu
   (vm)$ pip install tensorflow
   (vm)$ pip install tensorflow_datasets

2. Installa Flax:

   (vm)$ git clone https://github.com/google/flax.git
   (vm)$ pip install --user flax

3. Esegui lo script di addestramento Flax MNIST:

   (vm)$ cd flax/examples/mnist
   (vm)$ python3 main.py --workdir=/tmp/mnist \
       --config=configs/default.py \
       --config.learning_rate=0.05 \
       --config.num_epochs=5

Lo script scarica il set di dati e avvia l'addestramento. L'output dello script dovrebbe essere simile al seguente:

I0214 18:00:50.660087 140369022753856 train.py:146] epoch:  1, train_loss: 0.2421, train_accuracy: 92.97, test_loss: 0.0615, test_accuracy: 97.88
I0214 18:00:52.015867 140369022753856 train.py:146] epoch:  2, train_loss: 0.0594, train_accuracy: 98.16, test_loss: 0.0412, test_accuracy: 98.72
I0214 18:00:53.377511 140369022753856 train.py:146] epoch:  3, train_loss: 0.0418, train_accuracy: 98.72, test_loss: 0.0296, test_accuracy: 99.04
I0214 18:00:54.727168 140369022753856 train.py:146] epoch:  4, train_loss: 0.0305, train_accuracy: 99.06, test_loss: 0.0257, test_accuracy: 99.15
I0214 18:00:56.082807 140369022753856 train.py:146] epoch:  5, train_loss: 0.0252, train_accuracy: 99.20, test_loss: 0.0263, test_accuracy: 99.18