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 |
| 成像網 | 〜5分鐘 | imagenet-64x64.png |
探索不同採樣策略的最簡單方法是直接修改example.py 。您還可以通過簡單地複制相關位，將預訓練的模型和/或我們建議的EDM採樣器合併到您自己的代碼中。請注意，預訓練模型的類定義存儲在泡菜本身中，並在通過torch_utils.persistence中無釘時自動加載。要在外部Python腳本中使用這些模型，只需確保通過PYTHONPATH可以通過torch_utils和dnnlib進行輔助即可。
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) ：啟用所有GPU時，請與當前用戶的UID/GID一起運行交互式會話，以避免將Docker寫入root。-v `pwd`:/scratch --workdir /scratch ：安裝電流運行dir（例如，在主機機器上的GIT回購的頂部）to /scratch在容器中，並將其用作當前工作dir。-e HOME=/scratch ：指定在哪裡緩存臨時文件。注意：如果您想要更多的細粒控制，則可以設置DNNLIB_CACHE_DIR （用於預訓練的模型下載緩存）。您希望這些高速緩存DIR可以駐留在持久的捲上，以便它們的內容保留在多個docker run Inconcations中。 我們為我們建議的培訓配置（配置F）以及基線配置（config 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成立距離（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張生成的圖像；提供較少的圖像將導致錯誤，而提供更多圖像將使用隨機子集。為了減少隨機變化的效果，我們建議使用不同種子多次重複計算，例如--seeds=0-49999 ， --seeds=50000-99999 ，和--seeds=100000-149999 。在我們的論文中，我們計算了每個FID三次，並報告了最小值。
另請注意，將生成的圖像與最初對模型訓練的同一數據集進行比較很重要。為了促進評估，我們提供了與我們的預培訓模型相對應的確切參考統計信息：
對於ImagEnet，我們提供兩組參考統計信息，以啟用蘋果到應用的比較： imagenet-64x64.npz在評估EDM模型（ edm-imagenet-64x64-cond-adm.pkl ）時應使用，而應在評估基礎3時使用ImageNet 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：以1024x1024圖像下載Flickr-Faces-HQ數據集，並以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：下載更新的動物臉-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對象本地化挑戰，並以64x64分辨率轉換為ZIP存檔：
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上面的示例使用默認批處理大小為512張圖像（由--batch控制），該圖像均勻分配在8個GPU（由--nproc_per_node控制）中，每GPU產生64張圖像。訓練大型型號可能用盡GPU內存；避免這種情況的最佳方法是限制Per-GPU批處理大小，例如 - --batch-gpu=32 。這採用梯度積累來產生與使用全GPU批次相同的結果。有關選項的完整列表，請參見python train.py --help 。
每個訓練運行的結果都保存到新創建的目錄中，例如training-runs/00000-cifar10-cond-ddpmpp-edm-gpus8-batch64-fp32 。訓練環的導出網絡快照（ network-snapshot-*.pkl ）和培訓狀態（ training-state-*.pt ）定期（由--snap和--dump控制）。網絡快照可用於使用generate.py生成圖像，訓練狀態可用於以後（ --resume ）恢復培訓。其他有用的信息記錄在log.txt和stats.jsonl中。為了監視培訓融合，我們建議您查看訓練損失（ stats.jsonl中的"Loss/loss" ），並定期評估FID network-snapshot-*.pkl使用generate.py和fid.py
下表列出了我們用來獲取預訓練模型的確切培訓配置：
| 模型 | GPU | 時間 | 選項 |
|---|---|---|---|
| CIFAR10-32X32 -COND -VP | 8xv100 | 〜2天 | --cond=1 --arch=ddpmpp |
| CIFAR10-32X32 -COND -VE | 8xv100 | 〜2天 | --cond=1 --arch=ncsnpp |
| CIFAR10 −32x32 -uncond -vp | 8xv100 | 〜2天 | --cond=0 --arch=ddpmpp |
| CIFAR10-32X32 -UNCOND -VE | 8xv100 | 〜2天 | --cond=0 --arch=ncsnpp |
| FFHQ −64x64 -uncond -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 -uncond -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 -uncond -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，我們在四個NVIDIA DGX A100節點上進行了訓練，每個節點包含8個安培GPU，具有80 GB的內存。為了減少GPU內存需求，我們建議使用更多GPU訓練模型，或者使用-GPU限制-GPU的批量批量--batch-gpu 。要設置多節點培訓，請諮詢Torchrun文檔。
版權所有©2022，Nvidia Corporation＆Affiliates。版權所有。
所有材料，包括源代碼和預培訓模型，均在創意共享歸因於非商業期4.0國際許可下獲得許可。
baseline-cifar10-32x32-uncond-vp.pkl和baseline-cifar10-32x32-uncond-ve.pkl源自Yang Song，Yang Song，Jascha Sohl-Dickstein，DiedEderik P. Kingma，Abhishek Kumar，Stefano，Stefano Ermon和Ben Poole。這些模型最初是在Apache 2.0許可下共享的。
baseline-imagenet-64x64-cond-adm.pkl由Prafulla dhariwal和Alex Nichol衍生自訓練的模型。該模型最初是根據MIT許可證共享的。
imagenet-64x64-baseline.npz源自Prafulla Dhariwal和Alex Nichol的預定參考統計。統計數據最初是根據麻省理工學院許可證共享的。
@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，TuomasKynkäänniemi，Axel Sauer，Arash Vahdat和Janne Hellsten進行了討論和評論，以及Tero Kuosmanen，Samuel Klenberg和Janne Hellsten，用於維護我們的計算基礎架構。
