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 ）时应使用，而在评估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，用于维护我们的计算基础架构。
