EDM2
1.0.0
Menganalisis dan meningkatkan dinamika pelatihan model difusi
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine
https://arxiv.org/abs/2312.02696
Kontribusi utama:
conda env create -f environment.yml -n edmconda activate edmUntuk mereproduksi hasil utama dari makalah kami, cukup jalankan:
python example.pyIni adalah skrip mandiri minimal yang memuat model pra-terlatih terbaik untuk setiap dataset dan menghasilkan grid gambar 8x8 acak menggunakan pengaturan sampler optimal. Hasil yang diharapkan:
| Dataset | Runtime | Gambar referensi |
|---|---|---|
| CIFAR-10 | ~ 6 detik | cifar10-32x32.png |
| Ffhq | ~ 28 detik | ffhq-64x64.png |
| AFHQV2 | ~ 28 detik | afhqv2-64x64.png |
| Imagenet | ~ 5 menit | imagenet-64x64.png |
Cara termudah untuk mengeksplorasi berbagai strategi pengambilan sampel adalah dengan memodifikasi example.py secara langsung. Anda juga dapat memasukkan model pra-terlatih dan/atau sampler EDM yang diusulkan dalam kode Anda sendiri dengan cukup menyalin-mencelupkan bit yang relevan. Perhatikan bahwa definisi kelas untuk model pra-terlatih disimpan di dalam acar sendiri dan dimuat secara otomatis selama pembongkaran melalui torch_utils.persistence . Untuk menggunakan model dalam skrip Python eksternal, pastikan saja bahwa torch_utils dan dnnlib dapat diakses melalui PYTHONPATH .
Docker : Anda dapat menjalankan skrip contoh menggunakan Docker sebagai berikut:
# 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 Catatan: Gambar Docker memerlukan pelepasan driver NVIDIA r520 atau lebih baru.
Doa docker run mungkin terlihat menakutkan, jadi mari kita membongkar isinya di sini:
--gpus all -it --rm --user $(id -u):$(id -g) : Dengan semua GPU diaktifkan, jalankan sesi interaktif dengan UID/GID pengguna saat ini untuk menghindari file menulis Docker sebagai root.-v `pwd`:/scratch --workdir /scratch : pasang dir Dir saat ini (misalnya, bagian atas repo git ini pada mesin host Anda) untuk /scratch dalam wadah dan menggunakannya sebagai dir Dir yang berfungsi saat ini.-e HOME=/scratch : Tentukan di mana harus men -cache file sementara. Catatan: Jika Anda ingin lebih banyak kontrol berbutir halus, Anda dapat mengatur DNNLIB_CACHE_DIR (untuk cache unduhan model pra-terlatih). Anda ingin Dir Cache ini tinggal pada volume persisten sehingga isinya disimpan di beberapa doa docker run . Kami menyediakan model pra-terlatih untuk konfigurasi pelatihan yang kami usulkan (config f) serta konfigurasi baseline (config a):
Untuk menghasilkan batch gambar menggunakan model dan sampler yang diberikan, jalankan:
# 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 Menghasilkan sejumlah besar gambar dapat memakan waktu; Beban kerja dapat didistribusikan di beberapa GPU dengan meluncurkan perintah di atas menggunakan 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 Pengaturan sampler dapat dikontrol melalui opsi baris perintah; Lihat python generate.py --help untuk informasi lebih lanjut. Untuk hasil terbaik, kami sarankan menggunakan pengaturan berikut untuk setiap dataset:
# 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 Selain sampler EDM yang kami usulkan, generate.py juga dapat digunakan untuk mereproduksi ablasi sampler dari bagian 3 dari makalah kami. Misalnya:
# 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 Untuk menghitung jarak awal Fréchet (FID) untuk model dan sampler yang diberikan, pertama -tama menghasilkan 50.000 gambar acak dan kemudian membandingkannya dengan statistik referensi dataset menggunakan 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 Kedua perintah di atas dapat diparalelkan di beberapa GPU dengan menyesuaikan --nproc_per_node . Perintah kedua biasanya memakan waktu 1-3 menit dalam praktik, tetapi yang pertama kadang-kadang dapat memakan waktu beberapa jam, tergantung pada konfigurasi. Lihat python fid.py --help untuk daftar lengkap opsi.
Perhatikan bahwa nilai numerik FID bervariasi di berbagai biji acak dan sangat sensitif terhadap jumlah gambar. Secara default, fid.py akan selalu menggunakan 50.000 gambar yang dihasilkan; Memberikan lebih sedikit gambar akan menghasilkan kesalahan, sedangkan menyediakan lebih banyak akan menggunakan subset acak. Untuk mengurangi efek variasi acak, kami merekomendasikan mengulangi perhitungan beberapa kali dengan benih yang berbeda, misalnya, --seeds=0-49999 , --seeds=50000-99999 , dan --seeds=100000-149999 . Dalam makalah kami, kami menghitung setiap FID tiga kali dan melaporkan minimum.
Perhatikan juga bahwa penting untuk membandingkan gambar yang dihasilkan dengan dataset yang sama dengan model awalnya dilatih. Untuk memfasilitasi evaluasi, kami memberikan statistik referensi yang tepat yang sesuai dengan model pra-terlatih kami:
Untuk ImagEnet, kami menyediakan dua set statistik referensi untuk memungkinkan perbandingan apel-ke-apel: imagenet-64x64.npz harus digunakan saat mengevaluasi model EDM ( edm-imagenet-64x64-cond-adm.pkl ), sedangkan Model imagenet-64x64-baseline.npz harus digunakan ketika evalye harus digunakan (evalely the eveneling (evalely the eveneling (evalely the eveneling (evaleling (evalely the eveneling (eveneling-64x64-baseline. baseline-imagenet-64x64-cond-adm.pkl ); Yang terakhir ini awalnya dilatih oleh Dhariwal dan Nichol menggunakan data pelatihan yang sedikit berbeda.
Anda dapat menghitung statistik referensi untuk set data Anda sendiri sebagai berikut:
python fid.py ref --data=datasets/my-dataset.zip --dest=fid-refs/my-dataset.npz Dataset disimpan dalam format yang sama seperti di StyleGan: Arsip ZIP yang tidak terkompresi yang berisi file PNG yang tidak terkompresi dan data Metadata dataset.json untuk label. Dataset khusus dapat dibuat dari folder yang berisi gambar; Lihat python dataset_tool.py --help untuk informasi lebih lanjut.
CIFAR-10: Unduh versi CIFAR-10 Python dan konversi ke Zip Archive:
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: Unduh Dataset Flickr-Faces-HQ sebagai 1024x1024 gambar dan konversi ke ZIP Archive pada resolusi 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: Unduh Dataset Hewan Hewan-HQ yang diperbarui ( afhq-v2-dataset ) dan konversi ke ZIP Archive pada resolusi 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: Unduh Tantangan Lokalisasi Objek Imagenet dan Konversi ke Zip Archive pada resolusi 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 Anda dapat melatih model baru menggunakan train.py . Misalnya:
# 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 Contoh di atas menggunakan ukuran batch default dari 512 gambar (dikendalikan oleh --batch ) yang dibagi secara merata di antara 8 GPU (dikendalikan oleh --nproc_per_node ) untuk menghasilkan 64 gambar per GPU. Melatih model besar dapat kehabisan memori GPU; Cara terbaik untuk menghindari ini adalah dengan membatasi ukuran batch per-gpu, misalnya- --batch-gpu=32 . Ini menggunakan akumulasi gradien untuk menghasilkan hasil yang sama seperti menggunakan batch per-gpu penuh. Lihat python train.py --help untuk daftar lengkap opsi.
Hasil dari setiap pelatihan menjalankan disimpan ke direktori yang baru dibuat, misalnya training-runs/00000-cifar10-cond-ddpmpp-edm-gpus8-batch64-fp32 . Snapshot jaringan Ekspor Loop ( network-snapshot-*.pkl ) dan status pelatihan ( training-state-*.pt ) secara berkala (dikendalikan oleh --snap dan --dump ). Snapshot jaringan dapat digunakan untuk menghasilkan gambar dengan generate.py , dan status pelatihan dapat digunakan untuk melanjutkan pelatihan nanti ( --resume ). Informasi berguna lainnya dicatat dalam log.txt dan stats.jsonl . Untuk memantau konvergensi pelatihan, kami sarankan untuk melihat kehilangan pelatihan ( "Loss/loss" di stats.jsonl ) serta secara berkala mengevaluasi FID untuk network-snapshot-*.pkl menggunakan generate.py dan fid.py
Tabel berikut mencantumkan konfigurasi pelatihan yang tepat yang kami gunakan untuk mendapatkan model pra-terlatih kami:
| Model | GPU | Waktu | Opsi |
|---|---|---|---|
| CIFAR10-32X32 - COND -VP | 8xv100 | ~ 2 hari | --cond=1 --arch=ddpmpp |
| CIFAR10-32X32 - COND -VE | 8xv100 | ~ 2 hari | --cond=1 --arch=ncsnpp |
| CIFAR10-32X32 -ONCOND -VP | 8xv100 | ~ 2 hari | --cond=0 --arch=ddpmpp |
| CIFAR10-32X32 -ONCOND -VE | 8xv100 | ~ 2 hari | --cond=0 --arch=ncsnpp |
| FFHQ - 64X64 -ONCOND -VP | 8xv100 | ~ 4 hari | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| FFHQ - 64X64 -ONCOND -VE | 8xv100 | ~ 4 hari | --cond=0 --arch=ncsnpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| AFHQV2-64X64 -ONCOND -VP | 8xv100 | ~ 4 hari | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 |
| AFHQV2-64X64 -ONCOND -VE | 8xv100 | ~ 4 hari | --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 hari | --cond=1 --arch=adm --duration=2500 --batch=4096 --lr=1e-4 --ema=50 --dropout=0.10 --augment=0 --fp16=1 --ls=100 --tick=200 |
Untuk ImageNet-64, kami menjalankan pelatihan pada empat node NVIDIA DGX A100, masing-masing berisi 8 GPU Ampere dengan memori 80 GB. Untuk mengurangi persyaratan memori GPU, kami merekomendasikan pelatihan model dengan lebih banyak GPU atau membatasi ukuran batch per-GPU dengan --batch-gpu . Untuk mengatur pelatihan multi-node, silakan berkonsultasi dengan dokumentasi Torchrun.
Hak Cipta © 2022, Nvidia Corporation & Afiliasi. Semua hak dilindungi undang -undang.
Semua materi, termasuk kode sumber dan model pra-terlatih, dilisensikan di bawah Lisensi Internasional Atribusi Creative Commons-nonkomersial-Sharealike 4.0.
baseline-cifar10-32x32-uncond-vp.pkl dan baseline-cifar10-32x32-uncond-ve.pkl berasal dari model yang dilatih sebelumnya oleh Yang Lagu, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano, dan Benermon, dan Benermon. Model awalnya dibagikan di bawah lisensi Apache 2.0.
baseline-imagenet-64x64-cond-adm.pkl berasal dari model pra-terlatih oleh Prafulla Dhariwal dan Alex Nichol. Model ini awalnya dibagikan di bawah lisensi MIT.
imagenet-64x64-baseline.npz berasal dari statistik referensi yang telah dikomputasi oleh Prafulla Dhariwal dan Alex Nichol. Statistik awalnya dibagikan di bawah lisensi 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}
}
Ini adalah implementasi referensi penelitian dan diperlakukan sebagai penurunan kode satu kali. Karena itu, kami tidak menerima kontribusi kode luar dalam bentuk permintaan tarik.
Kami berterima kasih kepada Jaakko Lehtinen, Ming-Yu Liu, Tuomas Kynkäänniemi, Axel Sauer, Arash Vahdat, dan Janne Hellsten untuk diskusi dan komentar, dan tero Kuosmanen, Samuel Klenberg, dan Janne Hellsten karena mempertahankan infrastruktur komputur kami.
