Addestramento di modelli di diffusione con PyTorch

Questo tutorial mostra come addestrare modelli di diffusione sulle TPU utilizzando PyTorch Lightning e Pytorch XLA.

Obiettivi

  • Crea una Cloud TPU
  • Installa PyTorch Lightning
  • clona il repository di diffusione
  • Prepara il set di dati Imagenette
  • Esegui lo script di addestramento

Costi

In questo documento vengono utilizzati i seguenti componenti fatturabili di Google Cloud:

  • Compute Engine
  • Cloud TPU

Per generare una stima dei costi in base all'utilizzo previsto, utilizza il Calcolatore prezzi. I nuovi utenti di Google Cloud possono essere idonei a una prova senza costi aggiuntivi.

Prima di iniziare

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

  1. Accedi al tuo account Google Cloud. Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti gratuiti per l'esecuzione, il test e il deployment dei carichi di lavoro.

  2. Nella pagina del selettore di progetti della console Google Cloud, seleziona o crea un progetto Google Cloud.

    Vai al selettore progetti

  3. Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.

  4. Nella pagina del selettore di progetti della console Google Cloud, seleziona o crea un progetto Google Cloud.

    Vai al selettore progetti

  5. Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.

    6. Questa procedura dettagliata utilizza i componenti fatturabili di Google Cloud. Consulta la pagina dei prezzi di Cloud TPU per stimare i costi. Assicurati di pulire le risorse che hai creato quando hai finito di utilizzarle per evitare addebiti inutili.

Crea una Cloud TPU

Queste istruzioni funzionano su TPU sia con un host singolo che con più host. Questo tutorial utilizza un comando v4-128, ma funziona in modo simile su tutte le dimensioni dell'acceleratore.

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

export ZONE=us-central2-b
export PROJECT_ID=your-project-id
export ACCELERATOR_TYPE=v4-128
export RUNTIME_VERSION=tpu-ubuntu2204-base
export TPU_NAME=your_tpu_name

Creare una Cloud TPU. 

gcloud compute tpus tpu-vm create ${TPU_NAME} \
--zone=${ZONE} \
--accelerator-type=${ACCELERATOR_TYPE} \
--version=${RUNTIME_VERSION} \
--subnetwork=tpusubnet

Installa il software richiesto

  1. Installare i pacchetti richiesti insieme all'ultima release PyTorch/XLA v2.2.0.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone=us-central2-b \
--worker=all \
--command="sudo apt-get update -y && sudo apt-get install libgl1 -y
git clone https://github.com/pytorch-tpu/stable-diffusion.git
cd stable-diffusion
pip install -e .
pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/master.zip -U
pip install clip
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"

  2. Correggi i file sorgente in modo che siano compatibili con Torch 2.2 e versioni successive.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone=us-central2-b \
--worker=all \
--command="cd ~/stable-diffusion/
sed -i \'s/from torch._six import string_classes/string_classes = (str, bytes)/g\' src/taming-transformers/taming/data/utils.py
sed -i \'s/trainer_kwargs\\[\"callbacks\"\\]/# trainer_kwargs\\[\"callbacks\"\\]/g\' main_tpu.py"

  3. Scarica Imagenette (una versione più piccola del set di dati Imagenet) e spostalo nella directory appropriata. 

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="wget -nv https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz
tar -xf  imagenette2.tgz
mkdir -p ~/.cache/autoencoders/data/ILSVRC2012_train/data
mkdir -p ~/.cache/autoencoders/data/ILSVRC2012_validation/data
mv imagenette2/train/*  ~/.cache/autoencoders/data/ILSVRC2012_train/data
mv imagenette2/val/* ~/.cache/autoencoders/data/ILSVRC2012_validation/data"

  4. Scarica il modello preaddestrato della prima fase.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="cd ~/stable-diffusion/
wget -nv -O models/first_stage_models/vq-f8/model.zip https://ommer-lab.com/files/latent-diffusion/vq-f8.zip
cd  models/first_stage_models/vq-f8/
unzip -o model.zip"

Addestra il modello

Esegui l'addestramento con questo comando:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="python3 stable-diffusion/main_tpu.py --train --no-test --base=stable-diffusion/configs/latent-diffusion/cin-ldm-vq-f8-ss.yaml -- data.params.batch_size=32 lightning.trainer.max_epochs=5 model.params.first_stage_config.params.ckpt_path=stable-diffusion/models/first_stage_models/vq-f8/model.ckpt lightning.trainer.enable_checkpointing=False lightning.strategy.sync_module_states=False"

Esegui la pulizia

Esegui una pulizia per evitare addebiti non necessari al tuo account dopo aver utilizzato le risorse che hai creato:

Utilizza Google Cloud CLI per eliminare la risorsa Cloud TPU.

  $  gcloud compute tpus delete diffusion-tutorial --zone=us-central2-b

Passaggi successivi

Prova le colab di PyTorch: