用于一般噪声逆问题的扩散后验采样（ICLR 2023亮点论文）

摘要

在本工作中，我们通过近似后验采样，将扩散求解器扩展到高效处理一般的噪声（非）线性逆问题。有趣的是，由此产生的后验采样方案是扩散采样与流形约束梯度的混合版本，且不包含严格的测量一致性投影步骤，这使得它在噪声环境下相比先前的研究能够产生更理想的生成路径。

前置条件

• Python 3.8

• PyTorch 1.11.0

• CUDA 11.3.1

• nvidia-docker（如果在 Docker 容器中使用 GPU）

也可以使用较低版本的 CUDA 并搭配合适的 PyTorch 版本。

例如：CUDA 10.2 配合 PyTorch 1.7.0

快速入门

1) 克隆仓库

git clone https://github.com/DPS2022/diffusion-posterior-sampling

cd diffusion-posterior-sampling

2) 下载预训练检查点

从链接下载检查点 "ffhq_10m.pt"，并将其放入 ./models/ 目录下。

mkdir models
mv {DOWNLOAD_DIR}/ffqh_10m.pt ./models/

其中 {DOWNLOAD_DIR} 是您下载检查点的目录。

:speaker: Imagenet 的检查点已上传。

3) 设置环境

[选项 1] 本地环境设置

我们使用外部代码来实现运动模糊和非线性去模糊。

git clone https://github.com/VinAIResearch/blur-kernel-space-exploring bkse

git clone https://github.com/LeviBorodenko/motionblur motionblur

安装依赖项

conda create -n DPS python=3.8

conda activate DPS

pip install -r requirements.txt

pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113

[选项 2] 构建 Docker 镜像

在运行以下命令之前，请先安装 Docker 引擎、GPU 驱动程序以及相应的 CUDA 版本。

Dockerfile 已经包含了克隆外部代码的命令，因此无需再次手动克隆。

--gpus=all 参数是使用本地 GPU 设备所必需的（Docker >= 19.03）

docker build -t dps-docker:latest .

docker run -it --rm --gpus=all dps-docker

4) 推理

python3 sample_condition.py \
--model_config=configs/model_config.yaml \
--diffusion_config=configs/diffusion_config.yaml \
--task_config={TASK-CONFIG};

:speaker: 对于 Imagenet 数据集，请使用 configs/imagenet_model_config.yaml

可能的任务配置

# 线性逆问题
- configs/super_resolution_config.yaml
- configs/gaussian_deblur_config.yaml
- configs/motion_deblur_config.yaml
- configs/inpainting_config.yaml

# 非线性逆问题
- configs/nonlinear_deblur_config.yaml
- configs/phase_retrieval_config.yaml

任务配置的结构

您需要在 data.root 中指定数据目录。默认值为 ./data/samples，其中包含来自 FFHQ 验证集的三张示例图像。

conditioning:
    method: # 在 guided_diffusion/condition_methods.py 中查看候选方法
    params:
        scale: 0.5

data:
    name: ffhq
    root: ./data/samples/

measurement:
    operator:
        name: # 在 guided_diffusion/measurements.py 中查看候选算子

noise:
    name:   # 高斯或泊松
    sigma:  # 如果使用高斯噪声，请设置此参数。
    (rate:) # 如果使用泊松噪声，请设置此参数。

引用

如果您对我们的工作感兴趣，请考虑引用：

@inproceedings{
chung2023diffusion,
title={Diffusion Posterior Sampling for General Noisy Inverse Problems},
author={Hyungjin Chung and Jeongsol Kim and Michael Thompson Mccann and Marc Louis Klasky and Jong Chul Ye},
booktitle={The Eleventh International Conference on Learning Representations },
year={2023},
url={https://openreview.net/forum?id=OnD9zGAGT0k}
}

Diffusion Posterior Sampling 快速上手指南

本指南帮助开发者快速部署并运行 Diffusion Posterior Sampling (DPS)，用于解决通用的含噪逆问题（如超分辨率、去模糊、图像修复等）。

1. 环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.8
• PyTorch: 1.11.0 (需匹配 CUDA 版本)
• CUDA: 11.3.1 (也可使用其他版本，如 CUDA 10.2 配合 PyTorch 1.7.0)
• GPU: 推荐使用 NVIDIA GPU 以加速推理
• 可选: Docker (若使用容器化部署)

2. 安装步骤

第一步：克隆代码仓库

git clone https://github.com/DPS2022/diffusion-posterior-sampling
cd diffusion-posterior-sampling

第二步：下载预训练模型

从 Google Drive 下载 FFHQ 数据集的预训练检查点 ffhq_10m.pt。 (注：如需 ImageNet 模型，该链接中也已上传)

将下载的文件移动到项目目录：

mkdir models
mv {DOWNLOAD_DIR}/ffhq_10m.pt ./models/

请将 {DOWNLOAD_DIR} 替换为您实际下载文件的路径。

第三步：配置运行环境

您可以选择 本地环境 或 Docker 两种方式之一。

选项 A：本地环境设置 (推荐)

本项目依赖外部代码库处理运动模糊和非线性去模糊，需先克隆：

git clone https://github.com/VinAIResearch/blur-kernel-space-exploring bkse
git clone https://github.com/LeviBorodenko/motionblur motionblur

创建 Conda 环境并安装依赖（国内用户建议使用清华或阿里镜像源加速）：

conda create -n DPS python=3.8
conda activate DPS

# 使用国内镜像源安装 requirements
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装指定版本的 PyTorch (CUDA 11.3)
pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113

选项 B：Docker 环境设置

如果您已安装 Docker Engine 和 NVIDIA Container Toolkit，可直接构建镜像：

docker build -t dps-docker:latest .

# 运行容器 (--gpus=all 用于启用本地 GPU)
docker run -it --rm --gpus=all dps-docker

Dockerfile 内部已自动处理外部代码库的克隆，无需手动执行 git clone 步骤。

3. 基本使用

完成环境配置后，即可运行推理脚本。以下是最基础的运行命令示例：

python3 sample_condition.py \
--model_config=configs/model_config.yaml \
--diffusion_config=configs/diffusion_config.yaml \
--task_config=configs/super_resolution_config.yaml

任务配置说明

通过修改 --task_config 参数，可以切换不同的逆问题任务：

线性逆问题：

• 超分辨率: configs/super_resolution_config.yaml
• 高斯去模糊: configs/gaussian_deblur_config.yaml
• 运动去模糊: configs/motion_deblur_config.yaml
• 图像修复 (Inpainting): configs/inpainting_config.yaml

非线性逆问题：

• 非线性去模糊: configs/nonlinear_deblur_config.yaml
• 相位检索: configs/phase_retrieval_config.yaml

注意：如果使用 ImageNet 模型，请将 --model_config 改为 configs/imagenet_model_config.yaml。默认数据路径为 ./data/samples/，您可在对应的 yaml 配置文件中修改 data.root 指向您的图片目录。