EDM2
1.0.0
확산 모델의 훈련 역학 분석 및 개선
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine
https://arxiv.org/abs/2312.02696
주요 기여 :
conda env create -f environment.yml -n edmconda activate edm논문의 주요 결과를 재현하려면 간단히 실행하십시오.
python example.py이 스크립트는 각 데이터 세트에 가장 적합한 미리 훈련 된 모델을로드하고 최적의 샘플러 설정을 사용하여 임의의 8x8 그리드를 생성하는 최소한의 독립형 스크립트입니다. 예상 결과 :
| 데이터 세트 | 실행 시간 | 참조 이미지 |
|---|---|---|
| Cifar-10 | ~ 6 초 | cifar10-32x32.png |
| FFHQ | ~ 28 초 | ffhq-64x64.png |
| AFHQV2 | ~ 28 초 | afhqv2-64x64.png |
| imagenet | ~ 5 분 | imagenet-64x64.png |
다른 샘플링 전략을 탐색하는 가장 쉬운 방법은 example.py 직접 수정하는 것입니다. 또한 관련 비트를 간단히 복사하여 미리 훈련 된 모델 및/또는 제안 된 EDM 샘플러를 자신의 코드에 통합 할 수 있습니다. 미리 훈련 된 모델에 대한 클래스 정의는 피클 자체 내에 저장되며 torch_utils.persistence 를 통해 꼬임을 뽑는 동안 자동으로로드됩니다. 외부 Python 스크립트에서 모델을 사용하려면 torch_utils 와 dnnlib PYTHONPATH 통해 적합한 지 확인하십시오.
Docker : 다음과 같이 Docker를 사용하여 예제 스크립트를 실행할 수 있습니다.
# 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 참고 : Docker 이미지에는 NVIDIA 드라이버 릴리스 r520 이상이 필요합니다.
docker run 호출은 어려워 보일 수 있으므로 여기에서 내용을 풀겠습니다.
--gpus all -it --rm --user $(id -u):$(id -g) : 모든 gpus를 활성화하면 Docker가 루트로 쓰는 것을 피하기 위해 현재 사용자의 UID/gid와 대화식 세션을 실행합니다.-v `pwd`:/scratch --workdir /scratch : 컨테이너의 /스크래치로 /스크래치를 사용하여 현재 작동하는 Dir로 사용하는 전류 러닝 DIR (예 : 호스트 시스템 의이 git repo의 상단)을 /scratch 하십시오.-e HOME=/scratch : 임시 파일을 캐시 할 위치를 지정하십시오. 참고 :보다 세밀한 제어를 원한다면 대신 DNNLIB_CACHE_DIR (미리 훈련 된 모델 다운로드 캐시의 경우)을 설정할 수 있습니다. 이 캐시 디스크가 지속적인 볼륨에 상주하여 해당 내용이 여러 docker run 호출에 걸쳐 유지되도록합니다. 제안 된 교육 구성 (구성 F)과 기준선 구성 (구성 A)에 미리 훈련 된 모델을 제공합니다.
주어진 모델과 샘플러를 사용하여 일련의 이미지를 생성하려면 다음을 실행하십시오.
# 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 많은 이미지를 생성하는 데 시간이 많이 걸릴 수 있습니다. torchrun 사용하여 위의 명령을 시작하여 워크로드는 여러 GPU에 배포 할 수 있습니다.
# 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 샘플러 설정은 명령 줄 옵션을 통해 제어 할 수 있습니다. 자세한 내용은 python generate.py --help 참조하십시오. 최상의 결과를 얻으려면 각 데이터 세트에 대해 다음 설정을 사용하는 것이 좋습니다.
# 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 제안 된 EDM 샘플러 외에도 generate.py 논문의 섹션 3에서 샘플러 절제를 재현하는 데 사용될 수 있습니다. 예를 들어:
# 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 주어진 모델 및 샘플러의 FRéchet Inception Disting (FID)을 계산하려면 먼저 50,000 개의 랜덤 이미지를 생성 한 다음 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 위의 두 명령 모두 --nproc_per_node 조정하여 여러 GPU에서 병렬화 할 수 있습니다. 두 번째 명령은 일반적으로 실제로 1-3 분이 걸리지 만 첫 번째 명령은 구성에 따라 때때로 몇 시간이 걸릴 수 있습니다. 전체 옵션 목록은 python fid.py --help 참조하십시오.
FID의 숫자 값은 임의의 씨앗에 따라 다르며 이미지 수에 매우 민감합니다. 기본적으로 fid.py 항상 50,000 개의 생성 이미지를 사용합니다. 더 적은 이미지를 제공하면 오류가 발생하지만 더 많은 것을 제공하면 랜덤 하위 집합이 사용됩니다. 무작위 변동의 영향을 줄이려면 다른 시드로 여러 번 계산을 반복하는 것이 좋습니다. eg, --seeds=0-49999 , --seeds=50000-99999 및 --seeds=100000-149999 . 우리 논문에서, 우리는 각 FID를 세 번 계산하고 최소값을보고했습니다.
또한 생성 된 이미지를 모델과 원래 훈련 한 것과 동일한 데이터 세트와 비교하는 것이 중요합니다. 평가를 용이하게하기 위해 미리 훈련 된 모델에 해당하는 정확한 참조 통계를 제공합니다.
ImageNet의 경우 사과-apples 비교를 가능하게하기위한 두 가지 참조 통계 세트를 제공합니다. EDM 모델 ( edm-imagenet-64x64-cond-adm.pkl )을 평가할 때 imagenet-64x64.npz 를 사용해야합니다. imagenet-64x64-baseline.npz 는 기본 모델을 평가할 때 사용해야합니다. baseline-imagenet-64x64-cond-adm.pkl ); 후자는 원래 약간 다른 훈련 데이터를 사용하여 Dhariwal과 Nichol에 의해 교육을 받았습니다.
다음과 같이 자신의 데이터 세트의 참조 통계를 계산할 수 있습니다.
python fid.py ref --data=datasets/my-dataset.zip --dest=fid-refs/my-dataset.npz 데이터 세트는 Stylegan과 동일한 형식으로 저장됩니다. 비 압축되지 않은 PNG 파일을 포함하는 압축되지 않은 지퍼 아카이브 및 레이블 용 메타 데이터 파일 dataset.json . 사용자 정의 데이터 세트는 이미지가 포함 된 폴더에서 만들 수 있습니다. 자세한 내용은 python dataset_tool.py --help 참조하십시오.
CIFAR-10 : CIFAR-10 Python 버전을 다운로드하고 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 : Flickr-Faces-HQ 데이터 세트를 1024x1024 이미지로 다운로드하고 64x64 해상도로 Zip Archive로 변환하십시오.
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 : 업데이트 된 Animal Faces-HQ 데이터 세트 ( afhq-v2-dataset )를 다운로드하고 64x64 해상도로 Zip Archive로 변환하십시오.
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 : ImageNet Object 현지화 챌린지를 다운로드하고 64x64 해상도에서 Zip Archive로 변환하십시오.
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 train.py 사용하여 새로운 모델을 훈련시킬 수 있습니다. 예를 들어:
# 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 위의 예는 8 GPU ( --nproc_per_node 로 제어)로 균일하게 나뉘어 GPU 당 64 개의 이미지를 생성하는 512 개의 이미지 ( --batch 로 제어)의 기본 배치 크기를 사용합니다. 대형 모델 훈련은 GPU 메모리가 부족할 수 있습니다. 이를 피하는 가장 좋은 방법은 GPU 당 배치 크기 (예 : --batch-gpu=32 제한하는 것입니다. 이는 GPU 당 배치를 사용하는 것과 동일한 결과를 산출하기 위해 구배 축적을 사용합니다. 전체 옵션 목록은 python train.py --help 참조하십시오.
각 교육 실행의 결과는 새로 생성 된 디렉토리 (예 : training-runs/00000-cifar10-cond-ddpmpp-edm-gpus8-batch64-fp32 같은 디렉토리로 저장됩니다. 교육 루프는 정기적 인 간격 ( --snap 및 --dump 로 제어)으로 네트워크 스냅 샷 ( network-snapshot-*.pkl training-state-*.pt 네트워크 스냅 샷은 generate.py 와 함께 이미지를 생성하는 데 사용될 수 있으며 교육 상태를 사용하여 나중에 교육을 재개하는 데 ( --resume ). 다른 유용한 정보는 log.txt 및 stats.jsonl 에 기록됩니다. 교육 수렴을 모니터링하려면 stats.jsonl 의 교육 손실 ( "Loss/loss" )을 살펴보고 generate.py 및 fid.py 사용하여 network-snapshot-*.pkl 에 대한 FID를 정기적으로 평가하는 것이 좋습니다.
다음 표에는 미리 훈련 된 모델을 얻는 데 사용한 정확한 교육 구성이 나와 있습니다.
| 모델 | gpus | 시간 | 옵션 |
|---|---|---|---|
| CIFAR10-32X32 − Cond -VP | 8xv100 | ~ 2 일 | --cond=1 --arch=ddpmpp |
| CIFAR10-32X32 − Cond -Ve | 8xv100 | ~ 2 일 | --cond=1 --arch=ncsnpp |
| CIFAR10-32X32- vp | 8xv100 | ~ 2 일 | --cond=0 --arch=ddpmpp |
| CIFAR10-32X32 · 비 on -VE | 8xv100 | ~ 2 일 | --cond=0 --arch=ncsnpp |
| FFHQ -64X64- vp | 8xv100 | ~ 4 일 | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| FFHQ -64X64 · UNCOND -VE | 8xv100 | ~ 4 일 | --cond=0 --arch=ncsnpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| AFHQV2-64X64- vp | 8xv100 | ~ 4 일 | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 |
| AFHQV2-64X64 · VIND -VE | 8xv100 | ~ 4 일 | --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 일 | --cond=1 --arch=adm --duration=2500 --batch=4096 --lr=1e-4 --ema=50 --dropout=0.10 --augment=0 --fp16=1 --ls=100 --tick=200 |
imagenet-64의 경우, 우리는 8 개의 메모리를 갖는 8 개의 암페어 GPU를 포함하는 4 개의 NVIDIA DGX A100 노드에 대한 교육을 실행했습니다. GPU 메모리 요구 사항을 줄이려면 GPU가 더 많은 모델을 훈련 시키거나 -GPU 당 배치 크기를 --batch-gpu 로 제한하는 것이 좋습니다. 멀티 노드 교육을 설정하려면 Torchrun 문서를 참조하십시오.
Copyright © 2022, Nvidia Corporation & Affiliates. 모든 권리 보유.
소스 코드 및 미리 훈련 된 모델을 포함한 모든 자료는 Creative Commons Attribution-Noncommercial-Sharealike 4.0 International License에 따라 라이센스가 부여됩니다.
baseline-cifar10-32x32-uncond-vp.pkl 및 baseline-cifar10-32x32-uncond-ve.pkl 은 Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon 및 Ben Pole의 미리 훈련 된 모델에서 파생됩니다. 이 모델은 원래 Apache 2.0 라이센스에 따라 공유되었습니다.
baseline-imagenet-64x64-cond-adm.pkl Prafulla Dhariwal 및 Alex Nichol에 의해 미리 훈련 된 모델에서 파생됩니다. 이 모델은 원래 MIT 라이센스에 따라 공유되었습니다.
imagenet-64x64-baseline.npz 는 Prafulla Dhariwal과 Alex Nichol의 미리 계산 된 참조 통계에서 파생됩니다. 통계는 원래 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}
}
이것은 연구 참조 구현이며 일회성 코드 드롭으로 취급됩니다. 따라서 당사는 풀 요청의 형태로 외부 코드 기여를 수락하지 않습니다.
우리는 Jaakko Lehtinen, Ming-Yu Liu, Tuomas Kynkääänniemi, Axel Sauer, Arash Vahdat 및 Janne Hellsten에게 토론과 의견에 대해 감사하고 Tero Kuosmanen, Samuel Klenberg 및 Janne Hellsten에게 감사의 말을 전합니다.
