Palette: 图像到图像的扩散模型

论文 | 项目

简介

这是 Palette: 图像到图像的扩散模型 的非官方 PyTorch 实现，主要继承自其超分辨率版本 Image-Super-Resolution-via-Iterative-Refinement。代码模板来源于我的另一个种子项目：distributed-pytorch-template。

以下是一些与论文描述相符的实现细节：

• 我们采用了 Guided-Diffusion 中使用的 U-Net 架构，这显著提升了生成样本的质量。
• 我们在低分辨率特征（16×16）中使用了类似 vanilla DDPM 的注意力机制。
• 在 Palette 中，我们对 $\gamma$ 而不是 $t$ 进行编码，并通过仿射变换将其嵌入。
• 如 Palette 所述，在推理过程中我们将方差 $Σ_\theta(x_t, t)$ 固定为一个常数。

当前状态

代码

• 扩散模型流水线
• 训练/测试流程
• 保存/加载训练状态
• 日志记录/TensorBoard
• 多 GPU 训练 (DDP)
• EMA
• 指标（目前支持 FID 和 IS）
• 数据集（目前支持修复、去裁剪和上色）
• Google Colab 脚本 🌟（目前支持修复）

任务

我计划按顺序完成以下任务：

• 在 CelebaHQ 上进行修复🚀 (Google Colab)
• 在 Places2 with 128×128 centering mask 上进行修复🚀

由于时间和 GPU 资源的限制，后续实验尚不确定：

• Places2 上的去裁剪
• ImageNet 验证集上的上色

结果

DDPM 模型需要大量的计算资源，而我们目前只构建了几组示例模型来验证本文中的想法。

可视化结果

Celeba-HQ

使用 200 个 epoch 和 93 万次迭代的结果，展示了在 centering mask 和 irregular mask 下的前 100 个样本。

Places2 with 128×128 centering mask

使用 16 个 epoch 和 66 万次迭代的结果，展示了在 centering mask 下的几组精选样本。

Places2 上的去裁剪

使用 8 个 epoch 和 33 万次迭代的结果，展示了在 uncropping 下的几组精选样本。

指标

任务	数据集	EMA	FID(-)	IS(+)
使用 centering mask 进行修复	Celeba-HQ	False	5.7873	3.0705
使用 irregular mask 进行修复	Celeba-HQ	False	5.4026	3.1221

使用方法

环境

pip install -r requirements.txt

预训练模型

数据集	任务	迭代次数	GPUs×Days×Bs	URL
Celeba-HQ	修复	93 万	2×5×3	Google Drive
Places2	修复	66 万	4×8×10	Google Drive

Bs 表示每个 GPU 的样本数量。

数据准备

我们大部分数据来自 Kaggle，可能与官方版本略有不同，你也可以从官方网站下载。

• CelebA-HQ 重采样（256x256）Kaggle
• Places2 官方 | Places2 Kaggle
• ImageNet 官方

我们使用这些数据集的默认划分来进行训练和评估。所使用的文件列表可以在 Celeba-HQ 和 Places2 中找到。

准备好自己的数据后，你需要修改相应的配置文件以指向你的数据。以下是一个示例：

"which_dataset": {  // 使用参数导入指定数据集 
    "name": ["data.dataset", "InpaintDataset"], // 导入 Dataset() 类
    "args":{ // 初始化数据集的参数
    	"data_root": "你的数据路径",
    	"data_len": -1,
    	"mask_mode": "hybrid"
    } 
},

关于 dataloader 和 validation split 的更多选项也可以在配置文件的 datasets 部分找到。

继续训练/从断点恢复训练

1. 从给定链接下载检查点文件。
2. 将配置文件中的 resume_state 设置为之前检查点所在的目录。以下为例，该目录包含训练状态和保存的模型：

"path": { //设置各个部分的文件路径
	"resume_state": "experiments/inpainting_celebahq_220426_150122/checkpoint/100" 
},

3. 在 model.py 文件的 load_everything 函数中设置你的网络标签，默认为 Network。按照教程中的设置，优化器和模型将分别从 100.state 和 100_Network.pth 中加载。

netG_label = self.netG.__class__.__name__
self.load_network(network=self.netG, network_label=netG_label, strict=False)

4. 运行脚本：

python run.py -p train -c config/inpainting_celebahq.json

我们测试了 SR3 和 Guided Diffusion 中使用的 U-Net 主干网络，在当前实验中，Guided Diffusion 的表现更为稳健。关于 主干网络、损失函数 和 评估指标 的更多选择，可以在配置文件的 which_networks 部分找到。

测试

1. 按照“数据准备”部分的步骤，修改配置文件以指向你的数据。
2. 按照“继续训练”部分的步骤，设置你的模型路径。
3. 运行脚本：

python run.py -p test -c config/inpainting_celebahq.json

评估

1. 创建两个文件夹，分别用于保存真实图像和生成样本图像，且两者的文件名需要一一对应。
2. 运行脚本：

python eval.py -s [真实图像路径] -d [生成图像路径]

致谢

我们的工作基于以下理论研究：

• 去噪扩散概率模型
• Palette：图像到图像的扩散模型
• 扩散模型在图像合成上超越 GAN

同时，我们也受益于以下开源项目：

• openai/guided-diffusion
• LouisRouss/Diffusion-Based-Model-for-Colorization

Palette 图像到图像扩散模型快速上手指南

本指南基于非官方 PyTorch 实现，帮助开发者快速部署并运行 Palette: Image-to-Image Diffusion Models，主要用于图像修复（Inpainting）、去裁剪（Uncropping）和上色（Colorization）等任务。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.8 或更高版本
• GPU: 支持 CUDA 的 NVIDIA 显卡（训练需要多卡支持，推理单卡即可）
• 依赖库: PyTorch, torchvision 及其他深度学习常用库

提示: 国内用户建议使用清华源或阿里源加速 pip 安装，以提升下载速度。

安装步骤

1. 克隆项目代码

   git clone https://github.com/Janspiry/Palette-Image-to-Image-Diffusion-Models.git
   cd Palette-Image-to-Image-Diffusion-Models
   

2. 安装依赖 使用 pip 安装所需依赖包。国内用户可添加 -i 参数指定镜像源：

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

3. 准备预训练模型（可选） 如果仅需进行测试或推理，可从 Google Drive 下载预训练权重（如 Celeba-HQ 或 Places2 数据集模型），并将其放置在项目目录中。

   • Celeba-HQ Inpainting: 下载链接
   • Places2 Inpainting: 下载链接

4. 准备数据集 下载所需数据集（如 CelebA-HQ 或 Places2），并修改配置文件中的 data_root 指向本地数据路径。

   • CelebA-HQ (Kaggle): 下载地址
   • Places2 (Kaggle): 下载地址

基本使用

以下以 Celeba-HQ 图像修复任务 为例，展示最基础的测试流程。

1. 配置修改

打开 config/inpainting_celebahq.json（或对应任务的配置文件），确保以下两项设置正确：

• 数据路径: 在 which_dataset -> args 中修改 data_root 为您的数据存放目录。
• 模型路径: 在 path 中修改 resume_state 为预训练模型所在的文件夹路径（例如：experiments/inpainting_celebahq_220426_150122/checkpoint/100）。

2. 运行测试（推理）

执行以下命令生成修复后的图像：

python run.py -p test -c config/inpainting_celebahq.json

程序将读取配置中的输入图像，利用加载的扩散模型进行去噪推理，并将结果保存在配置的输出目录中。

3. （可选）开始训练

如果您拥有足够的 GPU 资源并希望从头训练或微调模型，请使用以下命令：

python run.py -p train -c config/inpainting_celebahq.json

注意: 训练过程耗时较长且对显存要求较高，建议根据 config 文件中的 which_networks 部分调整 backbone（推荐使用 Guided Diffusion 版本的 U-Net）以适应您的硬件条件。