edm2

分析和改进扩散模型的训练动态（EDM2） CVPR 2024论文的官方PyTorch实现

预览图

分析和改进扩散模型的训练动态 Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine https://arxiv.org/abs/2312.02696

摘要：扩散模型凭借其对大规模数据集的无与伦比的扩展能力，目前在数据驱动的图像合成领域占据主导地位。在本文中，我们识别并纠正了流行的ADM扩散模型架构中导致训练不均匀和低效的几个原因，同时不改变其高层结构。我们观察到在训练过程中网络激活和权重的幅度变化和不平衡现象失控，因此重新设计了网络层以在期望上保持激活、权重和更新幅度。我们发现系统地应用这一理念消除了观察到的漂移和不平衡，在相同的计算复杂度下产生了明显更好的网络。我们的改进将ImageNet-512合成的先前最佳FID记录从2.41提升到1.81，这是通过快速确定性采样实现的。 作为独立的贡献，我们提出了一种在训练结束后设置指数移动平均（EMA）参数的方法。这允许在不进行多次训练的成本下精确调整EMA长度，并揭示了它与网络架构、训练时间和引导之间的令人惊讶的相互作用。

如需商业咨询，请访问我们的网站并提交表单：NVIDIA研究许可

要求

支持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。

Docker

为方便起见，我们提供了一个包含所需依赖项的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：将当前运行目录（例如，主机上此git仓库的顶层目录）挂载到容器内的/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长度和引导强度选择为最小化FDDINOv2。支持以下预设：

edm2-img512-{xs|s|m|l|xl|xxl}-fid        # 表2，无CFG
edm2-img512-{xs|s|m|l|xl|xxl}-guid-fid   # 表2，带CFG
edm2-img512-{xs|s|m|l|xl|xxl}-dino       # 表5，无CFG
edm2-img512-{xs|s|m|l|xl|xxl}-guid-dino  # 表5，带CFG
edm2-img64-{s|m|l|xl}-fid                # 表3

每个预设都对应一组特定选项，指向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

换句话说，我们将使用939524 kimg和EMA长度0.015的XXL尺寸条件模型，并以1.7的引导强度使用2147483 kimg的XS尺寸无条件模型进行引导。更多详情，请参见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 和 FDDINOv2，我们首先需要生成 50,000 张随机图像。这在实践中可能非常耗时，所以将工作负载分配到多个 GPU 上是有意义的。这可以通过 torchrun 启动 generate_images.py 来实现:

# 使用 8 个 GPU 生成 50000 张图像并将其保存为 out/*/*.png
torchrun --standalone --nproc_per_node=8 generate_images.py \
    --preset=edm2-img512-s-guid-fid --outdir=out --subdirs --seeds=0-49999

另外，generate_images.py 也可以在计算集群中作为多 GPU 或多节点作业启动。只要集群环境为每个 GPU 生成一个单独的进程并填充必要的环境变量，这就应该可以直接使用。更多详细信息，请参阅 torchrun 文档。

生成 50,000 张图像后，可以使用 calculate_metrics.py 计算 FID 和 FDDINOv2:

# 为 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×220）训练图像保存一次快照，根据模型大小，每次训练运行对应 118-635 GB 的数据。在 https://nvlabs-fi-cdn.nvidia.com/edm2/raw-snapshots/ 中，我们改为以 64Mi 间隔提供快照，每次训练运行对应 15-79 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/

上述命令下载 64 个网络 pickle，每个 238 MB，总计 14.8 GB。下载完成后，可以使用 reconstruct_phema.py 重构新的 EMA 配置:

# 重构一个新的 EMA 配置，std=0.150
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：

从Kaggle下载ILSVRC2012数据存档并将其解压到某处，例如downloads/imagenet。
裁剪并调整图像大小以创建原始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

通过预训练的VAE编码器运行图像以创建相应的潜在数据集：

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

计算原始RGB数据集的参考统计数据，用于calculate_metrics.py：

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

训练新模型

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

# 使用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数量和可用的VRAM。避免这种情况的最佳方法是使用梯度累积来限制每个GPU的批量大小。在上述示例中，总批量大小为2048张图像，即每个GPU 256张，但我们通过指定--batch-gpu=32将其限制为每个GPU 32张。修改--batch-gpu是安全的，因为它与其他超参数没有交互，而修改总批量大小也需要调整学习率等参数。

默认情况下，训练脚本每128Ki（= 128 kibi = 128×2^10）个训练图像打印一次状态（由--status控制），每8Mi（= 8×2^20）个训练图像保存一次网络快照（由--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。

许可证

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

引用

@inproceedings{Karras2024edm2,
  title     = {Analyzing and Improving the Training Dynamics of Diffusion Models},
  author    = {Tero Karras and Miika Aittala and Jaakko Lehtinen and
               Janne Hellsten and Timo Aila and Samuli Laine},
  booktitle = {Proc. CVPR},
  year      = {2024},
}

开发

这是一个研究参考实现，被视为一次性代码发布。因此，我们不接受以拉取请求形式的外部代码贡献。

致谢

我们感谢Eric Chan、Qinsheng Zhang、Erik Härkönen、Tuomas Kynkäänniemi、Arash Vahdat、Ming-Yu Liu和David Luebke的讨论和评论，以及Tero Kuosmanen和Samuel Klenberg维护我们的计算基础设施。