EDM2
1.0.0
Analyse und Verbesserung der Trainingsdynamik von Diffusionsmodellen
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine
https://arxiv.org/abs/2312.02696
Schlüsselbeiträge:
conda env create -f environment.yml -n edmconda activate edmUm die Hauptergebnisse aus unserem Papier zu reproduzieren, rennen Sie einfach:
python example.pyDies ist ein minimales eigenständiges Skript, das das beste vorgebildete Modell für jeden Datensatz lädt und ein zufälliges 8x8-Bildergitter unter Verwendung der optimalen Sampler-Einstellungen erzeugt. Erwartete Ergebnisse:
| Datensatz | Laufzeit | Referenzbild |
|---|---|---|
| CIFAR-10 | ~ 6 Sek | cifar10-32x32.png |
| Ffhq | ~ 28 Sek | ffhq-64x64.png |
| AFHQV2 | ~ 28 Sek | afhqv2-64x64.png |
| Bildnische | ~ 5 min | imagenet-64x64.png |
Der einfachste Weg, verschiedene Stichprobenstrategien zu untersuchen, besteht darin, example.py zu ändern. Sie können auch die vorgebildeten Modelle und/oder unseren vorgeschlagenen EDM-Sampler in Ihren eigenen Code einbeziehen, indem Sie einfach die entsprechenden Bits kopieren. Beachten Sie, dass die Klassendefinitionen für die vorgebreiteten Modelle in den Gurken selbst gespeichert und während des Entzugs über torch_utils.persistence automatisch geladen werden. Um die Modelle in externen Python -Skripten zu verwenden, stellen Sie einfach sicher, dass torch_utils und dnnlib durch PYTHONPATH akademisch sind.
Docker : Sie können das Beispielskript mit Docker wie folgt ausführen:
# 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 HINWEIS: Das Docker -Bild erfordert NVIDIA -Treiber -Release r520 oder höher.
Der docker run -Aufruf sieht möglicherweise entmutigend aus. Lassen Sie uns also hier den Inhalt auspacken:
--gpus all -it --rm --user $(id -u):$(id -g) : Führen Sie bei allen aktivierten GPUs eine interaktive Sitzung mit UID/GID des aktuellen Benutzers aus, um zu vermeiden, dass Docker -Schreiben von Dateien als Root geschrieben wird.-v `pwd`:/scratch --workdir /scratch : MOUNT Current Running Dir (z. B. die Oberseite dieses Git -Repo auf Ihrem Host -Computer) zum /scratch im Container und verwenden Sie diese als aktuell funktionierende Dir.-e HOME=/scratch : Geben Sie an, wo temporäre Dateien zwischenstrahlen sind. Hinweis: Wenn Sie eine feinkörnige Steuerung wünschen, können Sie stattdessen DNNLIB_CACHE_DIR festlegen (für vorgebreitete Modell-Download-Cache). Sie möchten, dass diese Cache -Dires auf anhaltenden Volumina liegen, damit ihr Inhalt über mehrere docker run -Aufrufe aufbewahrt wird. Wir bieten vorgebildete Modelle für unsere vorgeschlagene Trainingskonfiguration (config f) sowie für die Basiskonfiguration (config a):
Um eine Reihe von Bildern mit einem bestimmten Modell und einem bestimmten Sampler zu erzeugen, rennen Sie:
# 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 Eine große Anzahl von Bildern kann zeitaufwändig sein. Die Workload kann über mehrere GPUs verteilt werden, indem der obige Befehl mit torchrun gestartet wird:
# 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 Die Sampler-Einstellungen können über Befehlszeilenoptionen gesteuert werden. Weitere Informationen finden Sie python generate.py --help . Für die besten Ergebnisse empfehlen wir die Verwendung der folgenden Einstellungen für jeden Datensatz:
# 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 Neben unserem vorgeschlagenen EDM -Sampler generate.py kann auch verwendet werden, um die Samplerablationen aus Abschnitt 3 unseres Papiers zu reproduzieren. Zum Beispiel:
# 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 Um Fréchet Inception Distanz (FID) für ein bestimmtes Modell und einen bestimmten Stichproben zu berechnen, generieren Sie zunächst 50.000 zufällige Bilder und vergleichen Sie sie dann mit der Datensatzreferenzstatistik mit 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 Beide oben genannten Befehle können durch Anpassung --nproc_per_node über mehrere GPUs parallelisiert werden. Der zweite Befehl dauert normalerweise 1-3 Minuten in der Praxis, aber der erste kann je nach Konfiguration manchmal mehrere Stunden dauern. Siehe python fid.py --help für die vollständige Liste der Optionen.
Beachten Sie, dass der numerische Wert von FID zwischen verschiedenen zufälligen Samen variiert und sehr empfindlich gegenüber der Anzahl der Bilder ist. Standardmäßig wird fid.py immer 50.000 erzeugte Bilder verwenden. Die Bereitstellung weniger Bilder führt zu einem Fehler, während die Bereitstellung von mehr zufälligen Teilmengen verwendet wird. Um die Auswirkung von zufälligen Variationen zu verringern, empfehlen wir --seeds=50000-99999 die Berechnung mehrmals mit unterschiedlichen Samen zu wiederholen, --seeds=100000-149999 --seeds=0-49999 In unserem Artikel haben wir jede FID dreimal berechnet und das Minimum gemeldet.
Beachten Sie auch, dass es wichtig ist, die generierten Bilder mit demselben Datensatz zu vergleichen, mit dem das Modell ursprünglich trainiert wurde. Um die Bewertung zu erleichtern, stellen wir die genauen Referenzstatistiken bereit, die unseren vorgebrachten Modellen entsprechen:
Für ImagEnet stellen wir zwei Sätze von Referenzstatistiken zur Verfügung, um den Vergleich von Äpfeln zu Applen zu aktivieren: imagenet-64x64.npz sollte bei der Bewertung des EDM-Modells ( edm-imagenet-64x64-cond-adm.pkl ) verwendet werden, wobei das Basenmodell imagenet-64x64-baseline.npz ) verwendet wird. baseline-imagenet-64x64-cond-adm.pkl ); Letzteres wurde ursprünglich von Dhariwal und Nichol mit leicht unterschiedlichen Trainingsdaten trainiert.
Sie können die Referenzstatistiken für Ihre eigenen Datensätze wie folgt berechnen:
python fid.py ref --data=datasets/my-dataset.zip --dest=fid-refs/my-dataset.npz Datensätze werden im selben Format wie in Stylegan: Unkomprimierte ZIP -Archive mit unkomprimierten PNG -Dateien und einem Metadata dataset.json für Etiketten gespeichert. Benutzerdefinierte Datensätze können aus einem Ordner erstellt werden, der Bilder enthält. Weitere Informationen finden Sie unter python dataset_tool.py --help .
CIFAR-10: Laden Sie die CIFAR-10 Python-Version herunter und konvertieren Sie sie in ZIP-Archiv:
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: Laden Sie den Flickr-Faces-HQ-Datensatz als 1024x1024-Bilder herunter und konvertieren Sie sie in das ZIP-Archiv unter 64 x 64 Auflösung:
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: Laden Sie den aktualisierten Tier-HQ-Datensatz ( afhq-v2-dataset ) herunter und konvertieren Sie sie unter 64x64-Auflösung in das ZIP-Archiv:
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: Laden Sie die Herausforderung der ImageNet -Objektlokalisierung herunter und konvertieren Sie sie in das ZIP -Archiv unter 64x64 Auflösung:
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 Sie können neue Modelle mit train.py trainieren. Zum Beispiel:
# 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 Das obige Beispiel verwendet die Standard -Stapelgröße von 512 Bildern (gesteuert durch --batch ), die gleichmäßig auf 8 GPUs (gesteuert von --nproc_per_node ) aufgeteilt wird, um 64 Bilder pro GPU zu ergeben. Das Training großer Modelle kann das GPU -Speicher mehr haben. Der beste Weg, dies zu vermeiden, besteht darin, die Größe pro GPU-Stapelgröße, z. B.-- --batch-gpu=32 zu begrenzen. Dadurch wird die Gradientenakkumulation verwendet, um die gleichen Ergebnisse zu erzielen wie die Verwendung von vollem GPU-Stapeln. Siehe python train.py --help für die vollständige Liste der Optionen.
Die Ergebnisse jedes Trainingslaufs werden in einem neu erstellten Verzeichnis gespeichert, z training-runs/00000-cifar10-cond-ddpmpp-edm-gpus8-batch64-fp32 Die Trainingsschleife exportiert Network Snapshots ( network-snapshot-*.pkl ) und Trainingszustände ( training-state-*.pt ) in regelmäßigen Abständen (kontrolliert von --snap und --dump ). Mit den Netzwerk -Snapshots können Bilder mit generate.py generiert werden, und die Trainingszustände können zum späteren Wiederaufzeigen des Trainings verwendet werden ( --resume ). Weitere nützliche Informationen werden in log.txt und stats.jsonl aufgezeichnet. Um die Konvergenz der Schulungen zu überwachen, empfehlen wir, den Schulungsverlust ( "Loss/loss" in stats.jsonl ) zu betrachten und FID für network-snapshot-*.pkl mit generate.py und fid.py regelmäßig zu bewerten.
In der folgenden Tabelle werden die genauen Trainingskonfigurationen aufgeführt, mit denen wir unsere vorgebauten Modelle erhalten haben:
| Modell | GPUS | Zeit | Optionen |
|---|---|---|---|
| CIFAR10-32x32 -cond -vp | 8xv100 | ~ 2 Tage | --cond=1 --arch=ddpmpp |
| CIFAR10-32x32 -Cond -vest | 8xv100 | ~ 2 Tage | --cond=1 --arch=ncsnpp |
| CIFAR10-32X32 -OCNOND -VP | 8xv100 | ~ 2 Tage | --cond=0 --arch=ddpmpp |
| CIFAR10-32X32 -OCNOND -VEE | 8xv100 | ~ 2 Tage | --cond=0 --arch=ncsnpp |
| FFHQ -64X64 -OCNOND -VP | 8xv100 | ~ 4 Tage | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| FFHQ -64X64 -OCNOND -VE | 8xv100 | ~ 4 Tage | --cond=0 --arch=ncsnpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| AFHQV2-64X64 -OCNOND -VP | 8xv100 | ~ 4 Tage | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 |
| AFHQV2-64X64 -OCNOND -VE | 8xv100 | ~ 4 Tage | --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 Tage | --cond=1 --arch=adm --duration=2500 --batch=4096 --lr=1e-4 --ema=50 --dropout=0.10 --augment=0 --fp16=1 --ls=100 --tick=200 |
Für ImagEnet-64 haben wir das Training auf vier Nvidia DGX A100-Knoten durchgeführt, die jeweils 8 Ampere-GPUs mit 80 GB Speicher enthielten. Um die Anforderungen an die GPU-Speicher zu verringern, empfehlen wir, das Modell entweder mit mehr GPUs zu trainieren oder die Größe pro GPU mit --batch-gpu zu begrenzen. Um ein Multi-Knoten-Training einzurichten, wenden Sie sich bitte an die Torchrun-Dokumentation.
Copyright © 2022, Nvidia Corporation & Affiliates. Alle Rechte vorbehalten.
Alle Materialien, einschließlich Quellcode und vorgebreiteten Modellen, ist unter der Kreativ-Commons-Attribution-Noncommercial-Sharealike 4.0 International Lizenz lizenziert.
baseline-cifar10-32x32-uncond-vp.pkl und baseline-cifar10-32x32-uncond-ve.pkl stammen von den vorgebildeten Modellen von Yang, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Emermon und Benn Poole. Die Modelle wurden ursprünglich unter der Apache 2.0 -Lizenz geteilt.
baseline-imagenet-64x64-cond-adm.pkl wird vom vorgebildeten Modell von Prafulla Dhariwal und Alex Nichol abgeleitet. Das Modell wurde ursprünglich unter der MIT -Lizenz geteilt.
imagenet-64x64-baseline.npz stammt aus den vorberechtigten Referenzstatistiken von Prafula Dhariwal und Alex Nichol. Die Statistiken wurden ursprünglich im Rahmen der MIT -Lizenz geteilt.
@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}
}
Dies ist eine Forschungsreferenzimplementierung und wird als einmaliger Codeabfall behandelt. Daher akzeptieren wir keine externen Codebeiträge in Form von Pull -Anfragen.
Wir danken Jaakko Lehtinen, Ming-yu Liu, Tuomas Kynkäänniemi, Axel Sauer, Arash Vahdat und Janne Hellsten für Diskussionen und Kommentare und tero Kuoshmen, Samuel Klenberg und Janne Hellsten für die Pflege unserer Zusammensetzung der Infrasstruktur.
