去噪扩散概率模型

去噪扩散概率模型的非官方 PyTorch 实现 [1]。

该实现尽可能地遵循了官方 TensorFlow 实现 [2] 的细节。我采用了 PyTorch 风格的代码编写方式，将 [2] 移植到 PyTorch 中，希望熟悉 PyTorch 的开发者能够轻松理解每个实现细节。

待办事项

• 数据集
  • 支持 CIFAR10
  • 支持 LSUN
  • 支持 CelebA-HQ
• 功能扩展
  • 梯度累积
  • 多 GPU 训练
• 实验复现
  • CIFAR10

环境要求

• Python 3.6

• 安装依赖包 为了安装最新版本的 TensorBoard，请先升级 pip：

  pip install -U pip setuptools
  pip install -r requirements.txt
  

• 下载数据集的预计算统计信息：

  cifar10.train.npz

  在项目目录下创建 stats 文件夹，并将 cifar10.train.npz 放入其中：

  stats
  └── cifar10.train.npz
  

从头开始训练

• 以 CIFAR10 为例：

  python main.py --train \
      --flagfile ./config/CIFAR10.txt
  

• [可选] 覆盖默认参数：

  python main.py --train \
      --flagfile ./config/CIFAR10.txt \
      --batch_size 64 \
      --logdir ./path/to/logdir
  

• [可选] 选择 GPU ID：

  CUDA_VISIBLE_DEVICES=1 python main.py --train \
      --flagfile ./config/CIFAR10.txt
  

• [可选] 多 GPU 训练：

  CUDA_VISIBLE_DEVICES=0,1,2,3 python main.py --train \
      --flagfile ./config/CIFAR10.txt \
      --parallel
  

评估

• 一个 flagfile.txt 文件会自动保存到你的日志目录中。对于 config/CIFAR10.txt，默认的日志目录是 ./logs/DDPM_CIFAR10_EPS。
• 开始评估：

  python main.py \
      --flagfile ./logs/DDPM_CIFAR10_EPS/flagfile.txt \
      --notrain \
      --eval
  

• [可选] 多 GPU 评估：

  CUDA_VISIBLE_DEVICES=0,1,2,3 python main.py \
      --flagfile ./logs/DDPM_CIFAR10_EPS/flagfile.txt \
      --notrain \
      --eval \
      --parallel
  

实验复现

CIFAR10

• FID: 3.249，Inception Score: 9.475(0.174)

检查点可以从我的 Google Drive 下载。

参考文献

[1] Denoising Diffusion Probabilistic Models

[2] 官方 TensorFlow 实现

pytorch-ddpm 快速上手指南

本指南基于非官方的 PyTorch 实现，帮助开发者快速复现去噪扩散概率模型（DDPM）。

环境准备

• 操作系统：Linux / macOS / Windows
• Python 版本：3.6 及以上
• 硬件要求：支持 CUDA 的 NVIDIA GPU（推荐用于训练和多卡加速）
• 前置依赖：
  • pip 和 setuptools（需更新至最新版以支持最新 TensorBoard）
  • 项目依赖包（见 requirements.txt）

安装步骤

1. 更新 pip 并安装依赖

   pip install -U pip setuptools
   pip install -r requirements.txt
   

   💡 国内加速建议：如遇下载缓慢，可添加清华或阿里镜像源：

   pip install -U pip setuptools -i https://pypi.tuna.tsinghua.edu.cn/simple
   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

2. 准备数据集统计文件

   • 下载预计算的统计文件：cifar10.train.npz
   • 在项目根目录创建 stats 文件夹，并将文件放入其中：

     stats
     └── cifar10.train.npz
     

   💡 注意：若无法访问 Google Drive，请自行寻找国内网盘搬运资源或通过科研协作渠道获取。

基本使用

以下以 CIFAR10 数据集为例，展示最基础的训练与评估流程。

1. 从头开始训练

使用默认配置启动训练：

python main.py --train \
    --flagfile ./config/CIFAR10.txt

可选参数示例：

• 自定义批次大小和日志路径：

  python main.py --train \
      --flagfile ./config/CIFAR10.txt \
      --batch_size 64 \
      --logdir ./path/to/logdir
  

• 指定单张 GPU：

  CUDA_VISIBLE_DEVICES=1 python main.py --train \
      --flagfile ./config/CIFAR10.txt
  

• 多卡并行训练：

  CUDA_VISIBLE_DEVICES=0,1,2,3 python main.py --train \
      --flagfile ./config/CIFAR10.txt \
      --parallel
  

2. 模型评估

训练完成后，系统会自动在日志目录（默认为 ./logs/DDPM_CIFAR10_EPS）保存 flagfile.txt。

启动评估：

python main.py \
    --flagfile ./logs/DDPM_CIFAR10_EPS/flagfile.txt \
    --notrain \
    --eval

多卡评估：

CUDA_VISIBLE_DEVICES=0,1,2,3 python main.py \
    --flagfile ./logs/DDPM_CIFAR10_EPS/flagfile.txt \
    --notrain \
    --eval \
    --parallel

📌 预训练模型：如需直接测试效果，可下载作者提供的 CIFAR10 检查点 放入对应日志目录后进行评估。