EDM2
1.0.0
Analisar e melhorar a dinâmica de treinamento dos modelos de difusão
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine
https://arxiv.org/abs/2312.02696
Principais contribuições:
conda env create -f environment.yml -n edmconda activate edmPara reproduzir os principais resultados do nosso artigo, basta executar:
python example.pyEste é um script independente mínimo que carrega o melhor modelo pré-treinado para cada conjunto de dados e gera uma grade de 8x8 aleatória de imagens usando as configurações ideais de amostrador. Resultados esperados:
| Conjunto de dados | Tempo de execução | Imagem de referência |
|---|---|---|
| Cifar-10 | ~ 6 seg | cifar10-32x32.png |
| Ffhq | ~ 28 seg | ffhq-64x64.png |
| AFHQV2 | ~ 28 seg | afhqv2-64x64.png |
| Imagenet | ~ 5 min | imagenet-64x64.png |
A maneira mais fácil de explorar diferentes estratégias de amostragem é modificar example.py diretamente. Você também pode incorporar os modelos pré-treinados e/ou nosso amostrador EDM proposto em seu próprio código, basta copiar os bits relevantes. Observe que as definições de classe para os modelos pré-treinadas são armazenadas dentro dos próprios picles e carregados automaticamente durante o desenrolar via torch_utils.persistence . Para usar os modelos em scripts de python externos, apenas verifique se torch_utils e dnnlib são acessíveis através PYTHONPATH .
Docker : Você pode executar o script de exemplo usando o Docker da seguinte maneira:
# Build the edm:latest image
docker build --tag edm:latest .
# Run the generate.py script using Docker:
docker run --gpus all -it --rm --user $( id -u ) : $( id -g )
-v ` pwd ` :/scratch --workdir /scratch -e HOME=/scratch
edm:latest
python example.py NOTA: A imagem do Docker requer liberação de driver NVIDIA r520 ou posterior.
A invocação docker run pode parecer assustadora, então vamos descompactar seu conteúdo aqui:
--gpus all -it --rm --user $(id -u):$(id -g) : com todas as GPUs ativadas, execute uma sessão interativa com o UID/GID do usuário atual para evitar que o docker escrevesse arquivos como root.-v `pwd`:/scratch --workdir /scratch : Monte o diretor de corrida atual (por exemplo, a parte superior deste repositório git na sua máquina host) para /scratch no contêiner e use -o como o diretor de trabalho atual.-e HOME=/scratch : especifique onde armazenar arquivos temporários. Nota: Se você deseja um controle mais refinado, pode definir DNNLIB_CACHE_DIR (para cache de download de modelo pré-treinado). Você deseja que esses diretores de cache residam em volumes persistentes, para que seu conteúdo seja retido em vários invocações docker run . Fornecemos modelos pré-treinados para nossa configuração de treinamento proposta (Config F), bem como a configuração da linha de base (Config A):
Para gerar um lote de imagens usando um determinado modelo e amostrador, execute:
# Generate 64 images and save them as out/*.png
python generate.py --outdir=out --seeds=0-63 --batch=64
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl Gerar um grande número de imagens pode ser demorado; A carga de trabalho pode ser distribuída por várias GPUs lançando o comando acima usando torchrun :
# Generate 1024 images using 2 GPUs
torchrun --standalone --nproc_per_node=2 generate.py --outdir=out --seeds=0-999 --batch=64
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl As configurações do amostrador podem ser controladas por meio de opções de linha de comando; Consulte python generate.py --help para obter mais informações. Para obter melhores resultados, recomendamos o uso das seguintes configurações para cada conjunto de dados:
# For CIFAR-10 at 32x32, use deterministic sampling with 18 steps (NFE = 35)
python generate.py --outdir=out --steps=18
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl
# For FFHQ and AFHQv2 at 64x64, use deterministic sampling with 40 steps (NFE = 79)
python generate.py --outdir=out --steps=40
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-ffhq-64x64-uncond-vp.pkl
# For ImageNet at 64x64, use stochastic sampling with 256 steps (NFE = 511)
python generate.py --outdir=out --steps=256 --S_churn=40 --S_min=0.05 --S_max=50 --S_noise=1.003
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-imagenet-64x64-cond-adm.pkl Além do nosso amostrador EDM proposto, generate.py também pode ser usado para reproduzir as ablações do amostrador da seção 3 do nosso artigo. Por exemplo:
# Figure 2a, "Our reimplementation"
python generate.py --outdir=out --steps=512 --solver=euler --disc=vp --schedule=vp --scaling=vp
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/baseline/baseline-cifar10-32x32-uncond-vp.pkl
# Figure 2a, "+ Heun & our {t_i}"
python generate.py --outdir=out --steps=128 --solver=heun --disc=edm --schedule=vp --scaling=vp
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/baseline/baseline-cifar10-32x32-uncond-vp.pkl
# Figure 2a, "+ Our sigma(t) & s(t)"
python generate.py --outdir=out --steps=18 --solver=heun --disc=edm --schedule=linear --scaling=none
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/baseline/baseline-cifar10-32x32-uncond-vp.pkl Para calcular a distância de FRÉCHET INCCECTION (FID) para um determinado modelo e amostrador, primeiro gere 50.000 imagens aleatórias e depois as compare com as estatísticas de referência do conjunto de dados usando fid.py :
# Generate 50000 images and save them as fid-tmp/*/*.png
torchrun --standalone --nproc_per_node=1 generate.py --outdir=fid-tmp --seeds=0-49999 --subdirs
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl
# Calculate FID
torchrun --standalone --nproc_per_node=1 fid.py calc --images=fid-tmp
--ref=https://nvlabs-fi-cdn.nvidia.com/edm/fid-refs/cifar10-32x32.npz Ambos os comandos acima podem ser paralelizados em várias GPUs ajustando --nproc_per_node . O segundo comando normalmente leva de 1 a 3 minutos na prática, mas o primeiro às vezes pode levar várias horas, dependendo da configuração. Consulte python fid.py --help para obter a lista completa de opções.
Observe que o valor numérico do FID varia entre diferentes sementes aleatórias e é altamente sensível ao número de imagens. Por padrão, fid.py sempre usará 50.000 imagens geradas; O fornecimento de menos imagens resultará em um erro, enquanto o fornecimento de mais usará um subconjunto aleatório. Para reduzir o efeito da variação aleatória, recomendamos repetir o cálculo várias vezes com diferentes sementes, por exemplo, --seeds=0-49999 , --seeds=50000-99999 e --seeds=100000-149999 . Em nosso artigo, calculamos cada FID três vezes e relatamos o mínimo.
Observe também que é importante comparar as imagens geradas com o mesmo conjunto de dados com o qual o modelo foi originalmente treinado. Para facilitar a avaliação, fornecemos as estatísticas de referência exatas que correspondem aos nossos modelos pré-treinados:
Para o ImageNet, fornecemos dois conjuntos de estatísticas de referência para ativar a comparação de maçãs para maçãs: imagenet-64x64.npz deve ser usado ao avaliar o modelo EDM ( edm-imagenet-64x64-cond-adm.pkl ), enquanto imagenet-64x64-baseline.npz deve ser usado baseline-imagenet-64x64-cond-adm.pkl ); Este último foi originalmente treinado por Dhariwal e Nichol usando dados de treinamento ligeiramente diferentes.
Você pode calcular as estatísticas de referência para seus próprios conjuntos de dados da seguinte forma:
python fid.py ref --data=datasets/my-dataset.zip --dest=fid-refs/my-dataset.npz Os conjuntos de dados são armazenados no mesmo formato que em Stylegan: Arquivos de Zip não compactados contendo arquivos PNG não compactados e um dataset.json de arquivos de metadados.json para etiquetas. Os conjuntos de dados personalizados podem ser criados a partir de uma pasta que contém imagens; Consulte python dataset_tool.py --help para obter mais informações.
CIFAR-10: Faça o download da versão CIFAR-10 Python e converta em Archive Zip:
python dataset_tool.py --source=downloads/cifar10/cifar-10-python.tar.gz
--dest=datasets/cifar10-32x32.zip
python fid.py ref --data=datasets/cifar10-32x32.zip --dest=fid-refs/cifar10-32x32.npzFFHQ: Faça o download do conjunto de dados do Flickr-Faces-HQ como imagens 1024x1024 e converta para o arquivo zip na resolução 64x64:
python dataset_tool.py --source=downloads/ffhq/images1024x1024
--dest=datasets/ffhq-64x64.zip --resolution=64x64
python fid.py ref --data=datasets/ffhq-64x64.zip --dest=fid-refs/ffhq-64x64.npz AFHQV2: Faça o download do conjunto de dados de Faces de Animal Atualizado ( afhq-v2-dataset ) e converta para Arquivo de Zip em resolução 64x64:
python dataset_tool.py --source=downloads/afhqv2
--dest=datasets/afhqv2-64x64.zip --resolution=64x64
python fid.py ref --data=datasets/afhqv2-64x64.zip --dest=fid-refs/afhqv2-64x64.npzImageNet: Baixe o desafio de localização de objetos ImageNet e converta em Arquivo de Zip na resolução 64x64:
python dataset_tool.py --source=downloads/imagenet/ILSVRC/Data/CLS-LOC/train
--dest=datasets/imagenet-64x64.zip --resolution=64x64 --transform=center-crop
python fid.py ref --data=datasets/imagenet-64x64.zip --dest=fid-refs/imagenet-64x64.npz Você pode treinar novos modelos usando train.py . Por exemplo:
# Train DDPM++ model for class-conditional CIFAR-10 using 8 GPUs
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-runs
--data=datasets/cifar10-32x32.zip --cond=1 --arch=ddpmpp O exemplo acima usa o tamanho do lote padrão de 512 imagens (controladas por --batch ) divididas uniformemente entre 8 GPUs (controlados por --nproc_per_node ) para produzir 64 imagens por GPU. Treinar grandes modelos pode ficar sem memória da GPU; A melhor maneira de evitar isso é limitar o tamanho do lote por GPU, por exemplo, --batch-gpu=32 . Isso emprega o acúmulo de gradiente para produzir os mesmos resultados que o uso de lotes completos por GPU. Consulte python train.py --help para obter a lista completa de opções.
Os resultados de cada execução de treinamento são salvos em um diretório recém-criado, por exemplo training-runs/00000-cifar10-cond-ddpmpp-edm-gpus8-batch64-fp32 . Os instantâneos de rede de exportações de loop de treinamento ( network-snapshot-*.pkl ) e estados de treinamento ( training-state-*.pt ) em intervalos regulares (controlados por --snap e --dump ). Os instantâneos de rede podem ser usados para gerar imagens com generate.py , e os estados de treinamento podem ser usados para retomar o treinamento posteriormente ( --resume ). Outras informações úteis são registradas em log.txt e stats.jsonl . Para monitorar a convergência do treinamento, recomendamos analisar a perda de treinamento ( "Loss/loss" em stats.jsonl ), bem como avaliar periodicamente o FID para network-snapshot-*.pkl usando generate.py e fid.py
A tabela a seguir lista as configurações exatas de treinamento que usamos para obter nossos modelos pré-treinados:
| Modelo | GPUs | Tempo | Opções |
|---|---|---|---|
| CIFAR10 - 32X32 -COND -VP | 8xv100 | ~ 2 dias | --cond=1 --arch=ddpmpp |
| Cifar10 - 32x32 -Cond -V | 8xv100 | ~ 2 dias | --cond=1 --arch=ncsnpp |
| CIFAR10 - 32X32 -UNCOND -VP | 8xv100 | ~ 2 dias | --cond=0 --arch=ddpmpp |
| CIFAR10 - 32X32 -Ucond -V | 8xv100 | ~ 2 dias | --cond=0 --arch=ncsnpp |
| FFHQ -64X64 -Ucond -VP | 8xv100 | ~ 4 dias | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| FFHQ -64X64 -Ucond -V | 8xv100 | ~ 4 dias | --cond=0 --arch=ncsnpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| AFHQV2 - 64X64 -Ucond -VP | 8xv100 | ~ 4 dias | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 |
| AFHQV2-64X64 -Ucond -V | 8xv100 | ~ 4 dias | --cond=0 --arch=ncsnpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 |
| ImageNet - 64x64 -Cond -Adm | 32XA100 | ~ 13 dias | --cond=1 --arch=adm --duration=2500 --batch=4096 --lr=1e-4 --ema=50 --dropout=0.10 --augment=0 --fp16=1 --ls=100 --tick=200 |
Para o ImageNet-64, executamos o treinamento em quatro nós NVIDIA DGX A100, cada um contendo GPUs de 8 amperes com 80 GB de memória. Para reduzir os requisitos de memória da GPU, recomendamos treinar o modelo com mais GPUs ou limitar o tamanho do lote por GPU com --batch-gpu . Para configurar o treinamento em vários nó, consulte a documentação da Torchrun.
Copyright © 2022, Nvidia Corporation & Affiliates. Todos os direitos reservados.
Todo o material, incluindo o código-fonte e os modelos pré-treinados, está licenciado sob a Licença Internacional Creative Commons Attribution-NonCommercial-Sharealike 4.0.
baseline-cifar10-32x32-uncond-vp.pkl e baseline-cifar10-32x32-uncond-ve.pkl são derivados dos modelos pré-treinados por Yang Song, Jascha Sohl-Dickstein, Diederik P. Os modelos foram originalmente compartilhados sob a licença Apache 2.0.
baseline-imagenet-64x64-cond-adm.pkl é derivado do modelo pré-treinado por Prafulla Dhariwal e Alex Nichol. O modelo foi originalmente compartilhado sob a licença do MIT.
imagenet-64x64-baseline.npz é derivado das estatísticas de referência pré-computadas de Prafulla Dhariwal e Alex Nichol. As estatísticas foram originalmente compartilhadas sob a licença do MIT.
@inproceedings{Karras2022edm,
author = {Tero Karras and Miika Aittala and Timo Aila and Samuli Laine},
title = {Elucidating the Design Space of Diffusion-Based Generative Models},
booktitle = {Proc. NeurIPS},
year = {2022}
}
Esta é uma implementação de referência de pesquisa e é tratada como uma queda de código única. Como tal, não aceitamos contribuições de código externas na forma de solicitações de tração.
Agradecemos a Jaakko Lehtinen, Ming-Yu Liu, Tuomas Kynkäänniemi, Axel Sauer, Arash Vahdat e Janne Hellsten por discussões e comentários, e Tero Kuosmanen, Samuel Klenberg e Janne Hellsten pela manutenção de nossa computa.
