EDM2与自动引导——PyTorch官方实现

分析并改进扩散模型的训练动态（CVPR 2024口头报告）
特罗·卡拉斯、米卡·艾塔拉、雅各·莱蒂宁、扬内·赫尔斯特恩、蒂莫·艾拉、萨穆利·莱内
https://arxiv.org/abs/2312.02696

用自身的一个较差版本引导扩散模型（NeurIPS 2024口头报告）
特罗·卡拉斯、米卡·艾塔拉、图奥马斯·金坎涅米、雅各·莱蒂宁、蒂莫·艾拉、萨穆利·莱内
https://arxiv.org/abs/2406.02507

如需商务合作，请访问我们的官网并提交表格：NVIDIA Research Licensing

系统要求

• 支持Linux和Windows，但出于性能和兼容性考虑，我们推荐使用Linux。
• 采样时需要1张及以上高端NVIDIA GPU，训练则需要8张及以上GPU。我们所有的测试和开发均在V100和A100 GPU上完成。
• 64位Python 3.9及PyTorch 2.1（或更高版本）。PyTorch安装说明请参见https://pytorch.org。
• 其他Python库：pip install click Pillow psutil requests scipy tqdm diffusers==0.26.3 accelerate==0.27.2
• 下载后期EMA重建所需的原始快照时，建议使用Rclone。

为方便起见，我们提供了一个包含所需依赖项的Dockerfile。使用方法如下：

# 构建Docker镜像
docker build --tag edm2:latest .

# 使用Docker运行generate_images.py
docker run --gpus all -it --rm --user $(id -u):$(id -g) \
    -v `pwd`:/scratch --workdir /scratch -e HOME=/scratch \
    edm2:latest \
    python generate_images.py --preset=edm2-img512-s-guid-dino --outdir=out

如果遇到错误，请确保已正确安装NVIDIA容器运行时。驱动程序兼容性详情请参阅NVIDIA PyTorch容器发布说明。

docker run命令行解析：

• --gpus all -it --rm --user $(id -u):$(id -g)：启用所有GPU，以当前用户的UID/GID运行交互式会话，避免Docker以root身份写入文件。
• -v pwd:/scratch --workdir /scratch：将当前工作目录（例如主机上的本仓库根目录）挂载到容器内的/scratch，并将其设为当前工作目录。
• -e HOME=/scratch：指定临时文件的缓存位置。若需更精细的控制，可设置DNNLIB_CACHE_DIR（用于预训练模型下载缓存）。这些缓存目录应位于持久化存储卷上，以便在多次docker run调用之间保留其内容。

使用预训练模型

我们提供了针对不同模型尺寸的EDM2配置（配置G）的预训练模型，分别在ImageNet-512和ImageNet-64数据集上训练。要使用给定模型生成图像，请运行：

# 生成几张图片并保存为out/*.png
python generate_images.py --preset=edm2-img512-s-guid-dino --outdir=out

上述命令会自动下载所需模型，并将其缓存在$HOME/.cache/dnnlib目录下，可通过设置DNNLIB_CACHE_DIR环境变量来覆盖该路径。--preset=edm2-img512-s-guid-dino选项表示我们将使用S尺寸的EDM2模型，在ImageNet-512数据集上训练，并采用指导采样方式，同时选择合适的EMA长度和指导强度以最小化FD_DINOv2。支持的预设如下：

# EDM2论文
edm2-img512-{xs|s|m|l|xl|xxl}-fid              # 表2，最小化fid
edm2-img512-{xs|s|m|l|xl|xxl}-dino             # 表5，最小化fd_dinov2
edm2-img64-{s|m|l|xl}-fid                      # 表3，最小化fid
edm2-img512-{xs|s|m|l|xl|xxl}-guid-{fid|dino}  # 表2，无分类器指导

# 自动引导论文
edm2-img512-{s|xxl}-autog-{fid|dino}           # 表1，条件式ImageNet-512
edm2-img512-s-uncond-autog-{fid|dino}          # 表1，无条件式ImageNet-512
edm2-img64-s-autog-{fid|dino}                  # 表1，条件式ImageNet-64

每个预设都对应一组特定的选项，指向https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions/中的模型。例如，--preset=edm2-img512-xxl-guid-dino等价于：

# 展开后的--preset=edm2-img512-xxl-guid-dino命令行
python generate_images.py \
    --net=https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions/edm2-img512-xxl-0939524-0.015.pkl \
    --gnet=https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions/edm2-img512-xs-uncond-2147483-0.015.pkl \
    --guidance=1.7 \
    --outdir=out

也就是说，我们将使用XXL尺寸的条件式模型，其kimg值为939524，EMA长度为0.015，并以XS尺寸的无条件式模型作为指导，该模型的kimg值为2147483，指导强度为1.7。更多细节请参阅generate_images.py中的config_presets。

计算FLOPs及指标

可使用count_flops.py估算给定模型的计算成本：

# 计算给定模型的FLOPs
python count_flops.py \
    https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions/edm2-img512-s-2147483-0.130.pkl

要计算FID和FD_DINOv2,首先需要生成50,000张随机图片。这在实际操作中可能非常耗时，因此最好将任务分配到多张GPU上执行。可以通过torchrun启动generate_images.py来实现：

# 使用8张GPU生成50,000张图片，并保存为out/*/*.png
torchrun --standalone --nproc_per_node=8 generate_images.py \
    --preset=edm2-img512-s-guid-fid --outdir=out --subdirs --seeds=0-49999

此外，也可以在计算集群中以多GPU或多节点作业的方式运行generate_images.py。只要集群环境能为每张GPU启动独立进程并设置好必要的环境变量，即可直接运行。更多详情请参阅torchrun文档。

生成50,000张图片后，即可使用calculate_metrics.py计算FID和FD_DINOv2：

# 计算 out/ 目录中随机子集50000张图像的指标
python calculate_metrics.py calc --images=out \
    --ref=https://nvlabs-fi-cdn.nvidia.com/edm2/dataset-refs/img512.pkl

此处，--ref 选项指向模型最初训练所用数据集的预计算参考统计信息。我们预训练模型所需的必要参考统计信息可在 https://nvlabs-fi-cdn.nvidia.com/edm2/dataset-refs/ 获取。

请注意，不同随机种子下指标的数值会有所不同，并且对图像数量非常敏感。默认情况下，calculate_metrics.py 使用 50,000 张生成图像，这符合既定的最佳实践。如果提供的图像数量少于该值，则会报错；而如果多于该值，则会使用一个随机子集。为了减少随机性带来的影响，我们建议使用不同的随机种子多次重复计算，例如 --seeds=0-49999、--seeds=50000-99999 和 --seeds=100000-149999。在我们的论文中，我们对每个指标进行了多次计算，并报告了最小值。

当需要对 EMA 长度或训练快照等进行更大规模的扫描时，按照上述方式使用 generate_images.py 可能不太现实。作为替代方案，也可以直接针对给定的网络 pickle 计算指标，同时即时生成所需的图像：

# 直接针对给定模型计算指标，无需保存任何图像
torchrun --standalone --nproc_per_node=8 calculate_metrics.py gen \
    --net=https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions/edm2-img512-s-2147483-0.130.pkl \
    --ref=https://nvlabs-fi-cdn.nvidia.com/edm2/dataset-refs/img512.pkl \
    --seed=123456789

我们还提供了必要的 API，以便从外部 Python 脚本以编程方式执行此类操作。有关详细信息，请参阅 calculate_metrics.py 中的 gen() 函数。

事后 EMA 重建

https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions/ 中的模型对应于特定的 EMA 长度选择。此外，我们还在 https://nvlabs-fi-cdn.nvidia.com/edm2/raw-snapshots/ 提供每次训练运行的原始快照，可用于重建任意 EMA 曲线。

需要注意的是，原始快照可能会占用大量磁盘空间。在论文中，我们每训练 8Mi（即 8 mebi = 8×2²⁰) 张图像保存一次快照，根据模型大小的不同，每次训练产生的数据量为 118–635 GB。而在 https://nvlabs-fi-cdn.nvidia.com/edm2/raw-snapshots/ 中，我们则以 32Mi 的间隔提供快照，每次训练产生的数据量为 30–159 GB。我们经过大量测试验证，确认这种频率足以实现准确的重建。

要重建新的 EMA 曲线，第一步是下载与特定训练运行相对应的原始快照。我们推荐使用 Rclone 来完成此操作：

# 下载预训练 edm2-img512-xs 模型的原始快照
rclone copy --progress --http-url https://nvlabs-fi-cdn.nvidia.com/edm2 \
    :http:raw-snapshots/edm2-img512-xs/ raw-snapshots/edm2-img512-xs/

上述命令将下载 128 个网络 pickle 文件，每个文件约 238 MB，总计 29.8 GB。下载完成后，可以使用 reconstruct_phema.py 重建新的 EMA 曲线：

# 重建 std=0.150 的新 EMA 曲线
python reconstruct_phema.py --indir=raw-snapshots/edm2-img512-xs \
    --outdir=out --outstd=0.150

该脚本会依次读取输入的每个 pickle 文件，并将重建后的模型保存到 out/phema-2147483-0.150.pkl，随后可与 generate_images.py 等工具配合使用。若需扫描不同的 EMA 长度，也可以同时重建多个 EMA 曲线：

# 同时重建 31 条 EMA 曲线，分 4 次流式处理输入数据
python reconstruct_phema.py --indir=raw-snapshots/edm2-img512-xs \
    --outdir=out --outstd=0.010,0.015,...,0.250 --batch=8

完整选项列表请参阅 python reconstruct_phema.py --help。

需要注意的是，我们的事后 EMA 方法并不局限于扩散模型，同样适用于其他类型的深度学习模型。要在您自己的训练中尝试这种方法，您可以：(1) 将 training/phema.py 添加到您的代码库中，(2) 修改您的训练循环以使用 phema.PowerFunctionEMA，以及 (3) 复制 reconstruct_phema.py，并根据您的需求进行调整。

数据集准备

数据集以未压缩的 ZIP 归档形式存储，其中包含未压缩的 PNG 或 NPY 文件，以及用于标注的元数据文件 dataset.json。在使用潜在扩散模型时，有必要为同一数据集创建两个版本：用于评估的原始 RGB 版本，以及用于训练的 VAE 编码的潜在版本。

以下是设置 ImageNet-512 的步骤：

1. 从 Kaggle 下载 ILSVRC2012 数据归档，并将其解压到某个位置，例如 downloads/imagenet。

2. 对图像进行裁剪和调整大小，以创建原始 RGB 数据集：

# 将原始 ImageNet 数据转换为 512x512 分辨率的 ZIP 归档
python dataset_tool.py convert --source=downloads/imagenet/ILSVRC/Data/CLS-LOC/train \
    --dest=datasets/img512.zip --resolution=512x512 --transform=center-crop-dhariwal

3. 将图像通过预训练的 VAE 编码器，以创建对应的潜在数据集：

# 将像素数据转换为 VAE 潜在表示
python dataset_tool.py encode --source=datasets/img512.zip \
    --dest=datasets/img512-sd.zip

4. 计算原始 RGB 数据集的参考统计信息，以便与 calculate_metrics.py 配合使用：

# 计算数据集参考统计信息，用于指标计算
python calculate_metrics.py ref --data=datasets/img512.zip \
    --dest=dataset-refs/img512.pkl

新模型训练

可以使用 train_edm2.py 训练新模型。例如，要使用与我们论文中相同的超参数训练一个 XS 尺寸的 ImageNet-512 条件模型，可运行：

# 使用8张GPU训练ImageNet-512的XS尺寸模型
torchrun --standalone --nproc_per_node=8 train_edm2.py \
    --outdir=training-runs/00000-edm2-img512-xs \
    --data=datasets/img512-sd.zip \
    --preset=edm2-img512-xs \
    --batch-gpu=32

本示例采用单节点、8张GPU进行训练，但在实际应用中，我们建议至少使用32张A100 GPU，即4个DGX节点。需要注意的是，训练大型模型时，根据GPU数量和显存容量的不同，很容易出现显存不足的情况。避免这一问题的最佳方法是通过梯度累积来限制每张GPU的批量大小。在上述示例中，总批量为2048张图像，即每张GPU 256张；但我们通过指定--batch-gpu=32将其限制为每张GPU 32张。调整--batch-gpu参数是安全的，因为它不会与其他超参数产生相互影响，而如果直接修改总批量，则还需要相应地调整学习率等其他超参数。

默认情况下，训练脚本每处理128Ki（= 128 kibi = 128×2¹⁰)张训练图像时打印一次状态信息（由--status控制），每8Mi（= 8×2²⁰)张训练图像时保存一次网络快照（由--snapshot控制），并每128Mi张训练图像时转储一次训练检查点（由--checkpoint控制）。状态信息会分别保存到log.txt（简要摘要）和stats.json（全面统计信息）文件中。网络快照则保存为network-snapshot-*.pkl格式，可以直接用于generate_images.py和reconstruct_phema.py等脚本。

训练检查点以training-state-*.pt格式保存，可用于稍后恢复训练。当训练脚本启动时，它会自动查找编号最高的检查点并加载（如果存在）。要恢复训练，只需再次运行相同的train_edm2.py命令行即可——务必使用完全相同的选项，以免在训练过程中意外更改超参数。若希望能够在任何时候暂停训练而不丢失进度，可以修改torch_utils/distributed.py中的should_suspend()函数，以实现所需的信号协议。

完整的选项列表请参阅python train_edm2.py --help。

二维玩具示例

自引导论文中使用的二维玩具示例可以通过toy_example.py重现：

# 使用自引导可视化采样分布。
python toy_example.py plot

完整的选项列表请参阅python toy_example.py --help。

许可证

版权所有 © 2024，NVIDIA公司及其关联公司。保留所有权利。

所有材料，包括源代码和预训练模型，均依据知识共享署名-非商业性使用-相同方式共享4.0国际许可协议授权。

引用

@inproceedings{Karras2024edm2,
  title     = {分析与改进扩散模型的训练动态},
  author    = {Tero Karras、Miika Aittala、Jaakko Lehtinen、Janne Hellsten、Timo Aila、Samuli Laine},
  booktitle = {CVPR会议论文集},
  year      = {2024},
}

@inproceedings{Karras2024autoguidance,
  title     = {用自身的一个较差版本引导扩散模型},
  author    = {Tero Karras、Miika Aittala、Tuomas Kynk\"a\"anniemi、Jaakko Lehtinen、Timo Aila、Samuli Laine},
  booktitle = {NeurIPS会议论文集},
  year      = {2024},
}

开发说明

本项目是一个研究参考实现，被视为一次性代码发布。因此，我们不接受来自外部的Pull Request形式的代码贡献。

致谢

我们感谢Eric Chan、Qinsheng Zhang、Erik Härkönen、Arash Vahdat、Ming-Yu Liu、David Luebke和Alex Keller提供的讨论与意见，同时也感谢Tero Kuosmanen和Samuel Klenberg对我们计算基础设施的维护工作。

EDM2 快速上手指南

EDM2 是 NVIDIA 研究团队推出的先进扩散模型训练与采样框架，包含两篇顶级会议论文成果（CVPR 2024 & NeurIPS 2024），旨在优化扩散模型的训练动态并引入“自引导”（Autoguidance）技术。

环境准备

系统要求

• 操作系统：支持 Linux 和 Windows，强烈推荐使用 Linux 以获得最佳性能和兼容性。
• GPU：
  • 图像生成（采样）：至少 1 张高端 NVIDIA GPU（测试基于 V100/A100）。
  • 模型训练：建议 8 张及以上 GPU。
• 软件版本：
  • Python 3.9 (64-bit)
  • PyTorch 2.1 或更高版本

依赖安装

确保已安装 PyTorch，然后安装其他必要的 Python 库：

pip install click Pillow psutil requests scipy tqdm diffusers==0.26.3 accelerate==0.27.2

提示：如需下载用于事后 EMA 重建的原始快照，建议安装 Rclone。

Docker 方案（推荐）

为避免环境配置冲突，项目提供了包含所有依赖的 Dockerfile。

1. 构建镜像：

   docker build --tag edm2:latest .
   

2. 运行生成脚本：

   docker run --gpus all -it --rm --user $(id -u):$(id -g) \
       -v `pwd`:/scratch --workdir /scratch -e HOME=/scratch \
       edm2:latest \
       python generate_images.py --preset=edm2-img512-s-guid-dino --outdir=out
   

   注意：请确保已正确安装 NVIDIA Container Runtime。

安装步骤

如果您选择直接在本地环境运行（非 Docker），请按以下步骤操作：

1. 克隆仓库：

   git clone https://github.com/NVlabs/edm2.git
   cd edm2
   

2. 安装 PyTorch： 访问 PyTorch 官网 根据您的 CUDA 版本安装对应的 PyTorch。国内用户可使用清华源加速：

   pip install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 安装项目依赖：

   pip install -r requirements.txt
   # 或者手动安装列出的库
   pip install click Pillow psutil requests scipy tqdm diffusers==0.26.3 accelerate==0.27.2
   

基本使用

EDM2 提供了预训练模型，可直接用于生成高质量图像。脚本会自动下载模型并缓存至 $HOME/.cache/dnnlib。

1. 生成图像（最简单示例）

使用预设配置生成图像并保存到 out 目录：

python generate_images.py --preset=edm2-img512-s-guid-dino --outdir=out

• --preset=edm2-img512-s-guid-dino：使用在 ImageNet-512 上训练的 S 尺寸模型，采用针对 FD_DINOv2 优化的引导策略。

2. 常用预设型号

您可以根据需求更换 --preset 参数：

• 最小化 FID (Table 2/3): edm2-img512-{xs|s|m|l|xl|xxl}-fid edm2-img64-{s|m|l|xl}-fid
• 最小化 FD_DINOv2 (Table 5): edm2-img512-{xs|s|m|l|xl|xxl}-dino
• 无分类器引导 (Classifier-free guidance): edm2-img512-{xs|s|m|l|xl|xxl}-guid-{fid|dino}
• 自引导 (Autoguidance) (NeurIPS 2024): edm2-img512-{s|xxl}-autog-{fid|dino} (条件生成) edm2-img512-s-uncond-autog-{fid|dino} (无条件生成)

3. 多 GPU 并行生成

为了评估指标（如 FID），通常需要生成 50,000 张图像。建议使用 torchrun 在多卡上分布式运行：

torchrun --standalone --nproc_per_node=8 generate_images.py \
    --preset=edm2-img512-s-guid-fid --outdir=out --subdirs --seeds=0-49999

4. 计算评估指标

生成图像后，使用以下命令计算 FID 和 FD_DINOv2：

python calculate_metrics.py calc --images=out \
    --ref=https://nvlabs-fi-cdn.nvidia.com/edm2/dataset-refs/img512.pkl

注意：指标数值对随机种子敏感，建议多次运行取最小值以获得稳定结果。