Addestramento di Resnet50 su Cloud TPU con PyTorch

Questo tutorial mostra come addestrare il modello ResNet-50 su un dispositivo Cloud TPU con PyTorch. Puoi applicare lo stesso pattern ad altri modelli di classificazione delle immagini ottimizzati per TPU che utilizzano PyTorch e il set di dati ImageNet.

Il modello in questo tutorial si basa sul Deep Residual Learning per il riconoscimento delle immagini, che introduce per la prima volta l'architettura della rete di residui (ResNet). Il tutorial utilizza la variante con 50 livelli, ResNet-50, e mostra l'addestramento del modello utilizzando PyTorch/XLA.

Obiettivi

  • Prepara il set di dati.
  • Esegui il job di addestramento.
  • Verifica i risultati dell'output.

Costi

In questo documento utilizzi i seguenti componenti fatturabili di Google Cloud:

  • Compute Engine
  • Cloud TPU

Per generare una stima dei costi basata sull'utilizzo previsto, utilizza il Calcolatore prezzi. I nuovi utenti di Google Cloud potrebbero essere idonei per una prova gratuita.

Prima di iniziare

Prima di iniziare questo tutorial, controlla che il progetto Google Cloud sia configurato correttamente.

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

  2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

    6. Questa procedura dettagliata utilizza i componenti fatturabili di Google Cloud. Controlla il Pagina dei prezzi di Cloud TPU per stimare i costi. Una volta terminato il loro utilizzo, assicurati di eseguire la pulizia delle risorse create per evitare addebiti superflui.

Crea una VM TPU

  1. Apri una finestra di Cloud Shell.

    Apri Cloud Shell

  2. Crea una VM TPU

    gcloud compute tpus tpu-vm create your-tpu-name \
--accelerator-type=v4-8 \
--version=tpu-ubuntu2204-base \
--zone=us-central2-b \
--project=your-project

  3. Connettiti alla VM TPU tramite SSH:

    gcloud compute tpus tpu-vm ssh  your-tpu-name --zone=us-central2-b

  4. Installa PyTorch/XLA sulla VM TPU:

    (vm)$ pip install torch~=2.4.0 torch_xla[tpu]~=2.4.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html

  5. Clona il repository GitHub di PyTorch/XLA

    (vm)$ git clone --depth=1 --branch r2.4 https://github.com/pytorch/xla.git

  6. Esegui lo script di addestramento con dati falsi

    (vm) $ PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=256 --num_epochs=1

Se riesci ad addestrare il modello utilizzando dati falsi, puoi provare ad addestrare il modello su dati reali, come ImageNet. Per istruzioni su come scaricare ImageNet, consulta Scaricare ImageNet. Nel comando dello script di addestramento, Il flag --datadir specifica la posizione del set di dati su cui eseguire l'addestramento. Il comando seguente presuppone che il set di dati ImageNet si trovi in ~/imagenet.

   (vm) $ PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py  --datadir=~/imagenet --batch_size=256 --num_epochs=1

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.

  1. Disconnettiti dalla VM TPU:

    (vm) $ exit

    Ora il tuo prompt dovrebbe essere username@projectname e mostrare in Cloud Shell.

  2. Elimina la VM TPU.

    $ gcloud compute tpus tpu-vm delete resnet50-tutorial \
   --zone=us-central2-b

Passaggi successivi