Neste tutorial, mostramos como treinar modelos de difusão em TPUs usando PyTorch Lightning e Pytorch XLA.
Objetivos
- Criar uma Cloud TPU
- Instalar o PyTorch Lightning
- Clonar o repositório de difusão
- Preparar o conjunto de dados do Imagenette
- execute o script de treinamento
Custos
Neste documento, você usará os seguintes componentes faturáveis do Google Cloud:
- Compute Engine
- Cloud TPU
Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços.
Antes de começar
Antes de começar o tutorial, verifique se o projeto do Google Cloud foi configurado corretamente.
- Faça login na sua conta do Google Cloud. Se você começou a usar o Google Cloud agora, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.
-
No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.
-
Verifique se a cobrança está ativada para o seu projeto do Google Cloud.
-
No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.
-
Verifique se a cobrança está ativada para o seu projeto do Google Cloud.
Este tutorial usa componentes faturáveis do Google Cloud. Consulte a página de preços da Cloud TPU para fazer uma estimativa dos custos. Para evitar cobranças desnecessárias, não se esqueça de apagar os recursos criados ao terminar de usá-los.
Criar uma Cloud TPU
Essas instruções funcionam em TPUs de host único e de vários hosts. Neste tutorial, usamos uma v4-128, mas ela funciona de maneira semelhante em todos os tamanhos de acelerador.
Configure algumas variáveis de ambiente para facilitar o uso dos comandos.
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
Criar um Cloud TPU.
gcloud compute tpus tpu-vm create ${TPU_NAME} \ --zone=${ZONE} \ --accelerator-type=${ACCELERATOR_TYPE} \ --version=${RUNTIME_VERSION} \ --subnetwork=tpusubnet
Instalar o software necessário
Instale os pacotes necessários com a versão mais recente do 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"
Os arquivos de origem foram corrigidos para serem compatíveis com a lanterna 2.2 e versões mais recentes.
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"
Faça o download do Imagenette, uma versão menor do conjunto de dados do Imagenet, e mova-o para o diretório apropriado.
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"
Faça o download do modelo pré-treinado de primeira etapa.
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"
Treine o modelo
Execute o treinamento com o seguinte 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"
Limpar
Execute uma limpeza para evitar cobranças desnecessárias na sua conta depois de usar os recursos criados:
Usar a CLI do Google Cloud para excluir o recurso da Cloud TPU.
$ gcloud compute tpus delete diffusion-tutorial --zone=us-central2-b
A seguir
Teste as colabs do PyTorch:
- Primeiros passos com o PyTorch em Cloud TPUs
- Como treinar MNIST em TPUs (em inglês)
- Como treinar o ResNet18 em TPUs com o conjunto de dados Cifar10 (em inglês)
- Inferência com modelo ResNet50 pré-treinado (em inglês)
- Transferência rápida de estilo neural
- Treinamento de vários núcleos da AlexNet no Fashion MNIST
- Treinamento de núcleo único da AlexNet no Fashion MNIST