去噪扩散修复模型 (DDRM)

arXiv | PDF | 项目官网

Bahjat Kawar¹, Michael Elad¹, Stefano Ermon², Jiaming Song^2
1 特克尼昂理工学院，² 斯坦福大学

DDRM 利用预训练的 DDPMs 来求解一般的线性逆问题。它能够高效地完成这一任务，且无需针对特定问题进行有监督训练。

运行实验

该代码已在 PyTorch 1.8 和 1.10 上测试通过。请参考 environment.yml 文件，了解可用于运行代码的 conda/mamba 环境列表。

预训练模型

我们使用来自以下仓库的预训练模型：https://github.com/openai/guided-diffusion、https://github.com/pesser/pytorch_diffusion 和 https://github.com/ermongroup/SDEdit。

为了与其他方法进行比较，我们使用了 ImageNet 验证集中的 1,000 张图像。这些图像的列表来自 https://github.com/XingangPan/deep-generative-prior/。

模型和数据集被放置在 exp/ 文件夹中，结构如下：

<exp> # 由 main.py 的 --exp 参数命名的文件夹
├── datasets # 所有数据集文件
│   ├── celeba # 所有 CelebA 文件
│   ├── imagenet # 所有 ImageNet 文件
│   ├── ood # 分布外的 ImageNet 图像
│   ├── ood_bedroom # 分布外的卧室图像
│   ├── ood_cat # 分布外的猫图像
│   └── ood_celeba # 分布外的 CelebA 图像
├── logs # 包含训练过程中生成的检查点和样本
│   ├── celeba
│   │   └── celeba_hq.ckpt # CelebA-HQ 的检查点文件
│   ├── diffusion_models_converted
│   │   └── ema_diffusion_lsun_<category>_model
│   │       └── model-x.ckpt # 第 x 次训练迭代保存的检查点文件
│   ├── imagenet # ImageNet 的检查点文件
│   │   ├── 256x256_classifier.pt
│   │   ├── 256x256_diffusion.pt
│   │   ├── 256x256_diffusion_uncond.pt
│   │   ├── 512x512_classifier.pt
│   │   └── 512x512_diffusion.pt
├── image_samples # 包含生成的样本
└── imagenet_val_1k.txt # ImageNet-1K 中使用的 1,000 张图像的列表。

需要注意的是，某些模型在无条件图像生成时可能无法生成高质量的样本，尤其是预训练的 CelebA 模型。

从模型中采样

从模型中采样的通用命令如下：

python main.py --ni --config {CONFIG}.yml --doc {DATASET} --timesteps {STEPS} --eta {ETA} --etaB {ETA_B} --deg {DEGRADATION} --sigma_0 {SIGMA_0} -i {IMAGE_FOLDER}

其中各选项含义如下：

• ETA 是论文中的 eta 超参数。（默认值：0.85）
• ETA_B 是论文中的 eta_b 超参数。（默认值：1）
• STEPS 控制采样过程中使用的步数。
• DEGREDATION 是允许的退化类型。（可选：cs2、cs4、inp、inp_lolcat、inp_lorem、deno、deblur_uni、deblur_gauss、deblur_aniso、sr2、sr4、sr8、sr16、sr_bicubic4、sr_bicubic8、sr_bicubic16、color）
• SIGMA_0 是 y 中观察到的噪声水平。
• CONFIG 是配置文件的名称（详见 configs/ 目录），包含批量大小、网络架构等超参数。
• DATASET 是所使用的数据集名称，用于确定检查点文件的位置。
• IMAGE_FOLDER 是存放生成图像的文件夹名称。（默认值：images）

例如，使用 ImageNet 256x256 无条件模型，对加噪图像进行 4 倍超分辨率重建，采用 20 步采样：

python main.py --ni --config imagenet_256.yml --doc imagenet --timesteps 20 --eta 0.85 --etaB 1 --deg sr4 --sigma_0 0.05

生成的图像将被放置在 <exp>/image_samples/{IMAGE_FOLDER} 文件夹中，其中 orig_{id}.png、y0_{id}.png、{id}_-1.png 分别对应原始图像、退化图像和恢复后的图像。

配置文件中包含一个设置，用于控制是否在训练数据集的分布内进行测试。

演示用图像

演示用图像列表可在以下链接找到：https://github.com/jiamings/ddrm-exp-datasets。将其放置在 <exp>/datasets 文件夹下，即可直接执行以下命令：

CelebA 加噪 4 倍超分辨率：

python main.py --ni --config celeba_hq.yml --doc celeba --timesteps 20 --eta 0.85 --etaB 1 --deg sr4 --sigma_0 0.05 -i celeba_hq_sr4_sigma_0.05

通用内容图像均匀去模糊：

python main.py --ni --config imagenet_256.yml --doc imagenet_ood --timesteps 20 --eta 0.85 --etaB 1 --deg deblur_uni --sigma_0 0.0 -i imagenet_sr4_sigma_0.0

卧室加噪 4 倍超分辨率：

python main.py --ni --config bedroom.yml --doc bedroom --timesteps 20 --eta 0.85 --etaB 1 --deg sr4 --sigma_0 0.05 -i bedroom_sr4_sigma_0.05

参考文献与致谢

@inproceedings{kawar2022denoising,
    title={Denoising Diffusion Restoration Models},
    author={Bahjat Kawar and Michael Elad and Stefano Ermon and Jiaming Song},
    booktitle={Advances in Neural Information Processing Systems},
    year={2022}
}

本实现基于或受以下项目启发：

• https://github.com/hojonathanho/diffusion（DDPM TensorFlow 仓库），
• https://github.com/pesser/pytorch_diffusion（用于加载 DDPM 模型的 PyTorch 工具），
• https://github.com/ermongroup/ddim（代码结构）。

DDRM 快速上手指南

去噪扩散恢复模型（DDRM）利用预训练的 DDPM 模型高效解决各类线性逆问题（如超分辨率、去模糊、修复等），无需针对特定问题进行监督训练。

环境准备

• 系统要求：Linux 或 macOS（Windows 需自行配置 PyTorch 环境）。
• Python 版本：推荐 Python 3.8+。
• 核心依赖：
  • PyTorch 1.8 或 1.10（已测试版本）。
  • Conda 或 Mamba（推荐用于管理虚拟环境）。
• 国内加速建议：
  • 安装 Conda 包时推荐使用清华源或中科大源。
  • 下载预训练模型若遇网络问题，可尝试使用镜像工具或手动下载后放入指定目录。

安装步骤

1. 克隆仓库

   git clone https://github.com/bahjat-kawar/ddrm.git
   cd ddrm
   

2. 创建并激活虚拟环境 使用项目提供的 environment.yml 创建环境（建议修改 yml 中的源为国内镜像以加速）：

   # 使用 conda
   conda env create -f environment.yml
   conda activate ddrm
   
   # 或使用 mamba (速度更快)
   mamba env create -f environment.yml
   mamba activate ddrm
   

3. 准备预训练模型与数据 代码依赖 OpenAI Guided Diffusion 等项目的预训练权重。请确保 exp/ 目录结构如下（需手动下载权重文件放入对应文件夹）：

   exp/
   ├── datasets/          # 放置数据集 (CelebA, ImageNet 等)
   ├── logs/              # 放置预训练权重 (.ckpt 或 .pt 文件)
   │   ├── celeba/
   │   ├── diffusion_models_converted/
   │   └── imagenet/
   └── imagenet_val_1k.txt
   

   注：具体模型下载地址请参考原文 README 中的 GitHub 链接。演示用图片数据集可从 ddrm-exp-datasets 获取并放入 exp/datasets。

基本使用

以下命令展示了如何使用预训练模型进行图像恢复。生成的图像将保存在 <exp>/image_samples/{IMAGE_FOLDER} 目录下。

通用命令格式

python main.py --ni --config {CONFIG}.yml --doc {DATASET} --timesteps {STEPS} --eta {ETA} --etaB {ETA_B} --deg {DEGRADATION} --sigma_0 {SIGMA_0} -i {IMAGE_FOLDER}

关键参数说明：

• --config: 配置文件名（位于 configs/ 目录），定义模型架构和批次大小。
• --doc: 数据集名称，用于定位对应的检查点文件。
• --timesteps: 采样步数（默认 20 步即可取得较好效果）。
• --deg: 退化类型，可选值包括：
  • 超分：sr2, sr4, sr8, sr16, sr_bicubic4 等
  • 去模糊：deblur_uni, deblur_gauss, deblur_aniso
  • 修复：inp, inp_lolcat, inp_lorem
  • 去噪：deno
  • 压缩感知：cs2, cs4
• --sigma_0: 观测噪声强度。

使用示例

示例 1：ImageNet 图像进行 4 倍超分辨率（带噪声） 使用 ImageNet 256x256 无条件模型，20 步采样，噪声标准差 0.05。

python main.py --ni --config imagenet_256.yml --doc imagenet --timesteps 20 --eta 0.85 --etaB 1 --deg sr4 --sigma_0 0.05

示例 2：CelebA 人脸图像进行 4 倍超分辨率 需先将演示数据放入 exp/datasets 目录。

python main.py --ni --config celeba_hq.yml --doc celeba --timesteps 20 --eta 0.85 --etaB 1 --deg sr4 --sigma_0 0.05 -i celeba_hq_sr4_sigma_0.05

示例 3：通用图像均匀去模糊

python main.py --ni --config imagenet_256.yml --doc imagenet_ood --timesteps 20 --eta 0.85 --etaB 1 --deg deblur_uni --sigma_0 0.0 -i imagenet_sr4_sigma_0.0

输出说明： 运行完成后，请在 exp/image_samples/{IMAGE_FOLDER} 中查看结果：

• orig_{id}.png: 原始清晰图像（Ground Truth，如有）。
• y0_{id}.png: 退化后的输入图像。
• {id}_-1.png: DDRM 恢复后的最终图像。