EDM2
1.0.0
Analyser et améliorer la dynamique de formation des modèles de diffusion
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine
https://arxiv.org/abs/2312.02696
Contributions clés:
conda env create -f environment.yml -n edmconda activate edmPour reproduire les principaux résultats de notre article, exécutez simplement:
python example.pyIl s'agit d'un script autonome minimal qui charge le meilleur modèle pré-formé pour chaque ensemble de données et génère une grille aléatoire 8x8 d'images en utilisant les paramètres d'échantillonneur optimaux. Résultats attendus:
| Ensemble de données | Temps d'exécution | Image de référence |
|---|---|---|
| Cifar-10 | ~ 6 sec | cifar10-32x32.png |
| Ffhq | ~ 28 sec | ffhq-64x64.png |
| Afhqv2 | ~ 28 sec | afhqv2-64x64.png |
| Imagenet | ~ 5 min | imagenet-64x64.png |
Le moyen le plus simple d'explorer différentes stratégies d'échantillonnage est de modifier directement example.py . Vous pouvez également intégrer les modèles pré-formés et / ou notre échantillonneur EDM proposé dans votre propre code en copier simplement les bits pertinents. Notez que les définitions de classe pour les modèles pré-formés sont stockées dans les cornichons eux-mêmes et chargés automatiquement pendant le décapage via torch_utils.persistence . Pour utiliser les modèles dans des scripts Python externes, assurez-vous simplement que torch_utils et dnnlib sont accessibles via PYTHONPATH .
Docker : vous pouvez exécuter l'exemple de script à l'aide de Docker comme suit:
# 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 Remarque: L'image Docker nécessite la libération du pilote NVIDIA r520 ou version ultérieure.
L'invocation docker run peut sembler intimidante, alors déballons son contenu ici:
--gpus all -it --rm --user $(id -u):$(id -g) : avec tous les GPU activés, exécutez une session interactive avec UID / GID de l'utilisateur actuel pour éviter les fichiers d'écriture docker sous forme de racine.-v `pwd`:/scratch --workdir /scratch : montez le répertoire en cours d'exécution (par exemple, le haut de ce repos git sur votre machine hôte) vers /scratch dans le conteneur et l'utiliser comme Dir de travail actuel.-e HOME=/scratch : spécifiez où mettre en cache des fichiers temporaires. Remarque: Si vous voulez un contrôle plus grainé, vous pouvez plutôt définir DNNLIB_CACHE_DIR (pour le cache de téléchargement du modèle pré-formé). Vous voulez que ces réseaux de cache résident sur des volumes persistants afin que leur contenu soit conservé à travers plusieurs invocations docker run . Nous fournissons des modèles pré-formés pour notre configuration de formation proposée (config F) ainsi que la configuration de base (config A):
Pour générer un lot d'images à l'aide d'un modèle donné et d'un échantillonneur, exécutez:
# 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 La génération d'un grand nombre d'images peut prendre du temps; La charge de travail peut être distribuée sur plusieurs GPU en lançant la commande ci-dessus à l'aide de 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 Les paramètres d'échantillonneur peuvent être contrôlés via des options de ligne de commande; Voir python generate.py --help pour plus d'informations. Pour les meilleurs résultats, nous vous recommandons d'utiliser les paramètres suivants pour chaque ensemble de données:
# 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 Outre notre échantillonneur EDM proposé, generate.py peut également être utilisé pour reproduire les ablations de l'échantillonneur de la section 3 de notre papier. Par exemple:
# 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 Pour calculer la distance de création de Fréchet (FID) pour un modèle et un échantillonneur donné, générez d'abord 50 000 images aléatoires, puis comparez-les avec les statistiques de référence de l'ensemble de données à l'aide de 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 Les deux commandes ci-dessus peuvent être parallélisées sur plusieurs GPU en ajustant --nproc_per_node . La deuxième commande prend généralement 1 à 3 minutes en pratique, mais la première peut parfois prendre plusieurs heures, selon la configuration. Voir python fid.py --help pour la liste complète des options.
Notez que la valeur numérique de la FID varie selon différentes graines aléatoires et est très sensible au nombre d'images. Par défaut, fid.py utilisera toujours 50 000 images générées; Fournir moins d'images entraînera une erreur, tandis que fournir davantage utilisera un sous-ensemble aléatoire. Pour réduire l'effet de la variation aléatoire, nous vous recommandons de répéter le calcul plusieurs fois avec différentes graines, par exemple, --seeds=0-49999 , --seeds=50000-99999 et --seeds=100000-149999 . Dans notre article, nous avons calculé chaque FID trois fois et signalé le minimum.
Notez également qu'il est important de comparer les images générées avec le même ensemble de données avec lequel le modèle a été initialement formé. Pour faciliter l'évaluation, nous fournissons les statistiques de référence exactes qui correspondent à nos modèles pré-formés:
Pour IMAMENET, nous fournissons deux ensembles de statistiques de référence pour permettre la comparaison des pommes à pommes: imagenet-64x64.npz doit être utilisée lors de l'évaluation du modèle EDM ( edm-imagenet-64x64-cond-adm.pkl ), tandis que imagenet-64x64-baseline.npz doit être utilisé lors de l'évaluation du modèle de base (Baseline (Baseline baseline-imagenet-64x64-cond-adm.pkl ); Ce dernier a été initialement formé par Dhariwal et Nichol en utilisant des données de formation légèrement différentes.
Vous pouvez calculer les statistiques de référence pour vos propres ensembles de données comme suit:
python fid.py ref --data=datasets/my-dataset.zip --dest=fid-refs/my-dataset.npz Les ensembles de données sont stockés dans le même format que In Stylegan: Archives ZIP non compressées contenant des fichiers PNG non compressés et un fichier de métadonnées dataset.json pour les étiquettes. Des ensembles de données personnalisés peuvent être créés à partir d'un dossier contenant des images; Voir python dataset_tool.py --help pour plus d'informations.
CIFAR-10: Téléchargez la version CIFAR-10 Python et convertissez en 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: Téléchargez l'ensemble de données Flickr-Faces-HQ sous forme d'images 1024x1024 et convertissez en archive zip à la résolution 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: Téléchargez l'ensemble de données Animal Faces-HQ mis à jour ( afhq-v2-dataset ) et converti en archive zip à la résolution 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: Téléchargez le défi de localisation des objets ImageNet et convertissez en archive zip à la résolution 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 Vous pouvez former de nouveaux modèles en utilisant train.py . Par exemple:
# 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 L'exemple ci-dessus utilise la taille du lot par défaut de 512 images (contrôlée par --batch ) qui est divisée uniformément entre 8 GPU (contrôlée par --nproc_per_node ) pour produire 64 images par GPU. La formation de grands modèles peut manquer de mémoire GPU; La meilleure façon d'éviter cela est de limiter la taille du lot Per-GPU, par exemple, --batch-gpu=32 . Cela utilise une accumulation de gradient pour donner les mêmes résultats que l'utilisation de lots Per-GPU. Voir python train.py --help pour la liste complète des options.
Les résultats de chaque course d'entraînement sont enregistrés dans un répertoire nouvellement créé, par exemple training-runs/00000-cifar10-cond-ddpmpp-edm-gpus8-batch64-fp32 . La boucle de formation exporte des instantanés de réseau ( network-snapshot-*.pkl ) et des états de formation ( training-state-*.pt ) à intervalles réguliers (contrôlé par --snap et --dump ). Les instantanés de réseau peuvent être utilisés pour générer des images avec generate.py , et les états de formation peuvent être utilisés pour reprendre la formation plus tard ( --resume ). D'autres informations utiles sont enregistrées dans log.txt et stats.jsonl . Pour surveiller la convergence de la formation, nous vous recommandons d'examiner la perte de formation ( "Loss/loss" dans stats.jsonl ) ainsi que de l'évaluation périodique du FID pour network-snapshot-*.pkl à l'aide de generate.py et fid.py
Le tableau suivant répertorie les configurations de formation exactes que nous avons utilisées pour obtenir nos modèles pré-formés:
| Modèle | GPUS | Temps | Options |
|---|---|---|---|
| CIFAR10‑32X32-COND - VP | 8xv100 | ~ 2 jours | --cond=1 --arch=ddpmpp |
| CIFAR10‑32X32 - VE | 8xv100 | ~ 2 jours | --cond=1 --arch=ncsnpp |
| CIFAR10‑32X32-ANCOND-VP | 8xv100 | ~ 2 jours | --cond=0 --arch=ddpmpp |
| CIFAR10‑32X32-ANCOND-VE | 8xv100 | ~ 2 jours | --cond=0 --arch=ncsnpp |
| FFHQ - 64X64-ANCOND-VP | 8xv100 | ~ 4 jours | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| FFHQ - 64X64-ANCOND-VE | 8xv100 | ~ 4 jours | --cond=0 --arch=ncsnpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| afhqv2‑64x64-ancond - vp | 8xv100 | ~ 4 jours | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 |
| AFHQV2‑64X64-ANCOND-VE | 8xv100 | ~ 4 jours | --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 jours | --cond=1 --arch=adm --duration=2500 --batch=4096 --lr=1e-4 --ema=50 --dropout=0.10 --augment=0 --fp16=1 --ls=100 --tick=200 |
Pour ImageNet-64, nous avons organisé la formation sur quatre nœuds NVIDIA DGX A100, chacun contenant 8 GPU d'ampère avec 80 Go de mémoire. Pour réduire les exigences de la mémoire du GPU, nous vous recommandons de former le modèle avec plus de GPU ou de limiter la taille du lot Per-GPU avec --batch-gpu . Pour mettre en place une formation multi-nœuds, veuillez consulter la documentation Torchrun.
Copyright © 2022, Nvidia Corporation & Affiliates. Tous droits réservés.
Tous les documents, y compris le code source et les modèles pré-formés, sont sous licence Creative Commons Attribution-NonCommercial-Sharealike 4.0 International Licence.
baseline-cifar10-32x32-uncond-vp.pkl et baseline-cifar10-32x32-uncond-ve.pkl sont dérivés des modèles pré-formés de Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon et Ben Poole. Les modèles ont été à l'origine partagés sous la licence Apache 2.0.
baseline-imagenet-64x64-cond-adm.pkl est dérivé du modèle pré-formé par Prafulla Dhariwal et Alex Nichol. Le modèle a été à l'origine partagé sous la licence du MIT.
imagenet-64x64-baseline.npz est dérivé des statistiques de référence précomputées par Prafulla Dhariwal et Alex Nichol. Les statistiques ont été à l'origine partagées sous la licence du 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}
}
Il s'agit d'une implémentation de référence de recherche et est traitée comme une chute de code ponctuelle. En tant que tels, nous n'acceptons pas les contributions du code extérieur sous forme de demandes de traction.
Nous remercions Jaakko Lehtinen, Ming-Yu Liu, Tuomas Kynkäänniemi, Axel Sauer, Arash Vahdat et Janne Hellsten pour des discussions et des commentaires, et Tero Kuosmanen, Samuel Klenberg et Janne Hellsten pour le maintien de notre infrastructure de calcul.
