扩散模型diffusion model用于图像恢复任务详细原理 (去雨，去雾等皆可)，附实现代码

话不多说，先上代码： 扩散模型diffusion model用于图像恢复完整可运行代码，附详细实验操作流程
令外一篇简化超分扩散模型SR3来实现图像恢复的博客见： 超分扩散模型 SR3 可以做图像去雨、去雾等恢复任务吗？

1. 去噪扩散概率模型

扩散模型是一类生成模型, 和生成对抗网络GAN 、变分自动编码器VAE和标准化流模型NFM等生成网络不同的是, 扩散模型在前向扩散过程中对图像逐步施加噪声, 直至图像被破坏变成完全的高斯噪声, 然后在反向采样过程中学习从高斯噪声还原为真实图像。在模型训练完成后，只需要随机给定一个高斯噪声，就可以生成丰富的真实图像。

2. 前向扩散

前向扩散过程就是向图像不断加高斯噪声，使其逐渐接近一个与输入数据相关的高斯分布。此处将未加噪声的数据记为$x_0$，$x_0\sim q(x_0)$，$q(x_0)$是为被噪声破坏的原始数据分布，则在$t$时刻的噪化状态和上一时刻$t-1$之间的关系为：
$$q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}\cdot x_{t-1},{\beta }_t\cdot \text{I}),\text{\hspace{1em}(1)}$$其中： t ∈ { 0 , 1 , . . . , T } t\in{\{0, 1, ..., T\}} t∈{0,1,...,T}，$\mathcal{N}$表示高斯噪声分布，${\beta }_t$是与时刻t相关的噪声方差调节因子，$\text{I}$是一个与初始状态$x_0$维度相同的单位矩阵。则输入$x_0$的条件下，$x_1,x_2,...,x_T$的联合分布可以表示为：
$$q(x_1,x_2,...,x_T|x_0)=\prod _{t=1}^{T}q(x_t|x_{t-1})\text{\hspace{1em}(2)}$$则根据根据马尔科夫性可以直接得到输入$x_0$的条件下$t$时刻的噪化状态为
$$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}\cdot x_0,(1-{\overline{\alpha }}_t)\cdot \text{I}),\text{\hspace{1em}(3)}$$其中：${\alpha }_t:=1-{\beta }_t$, ${\overline{\alpha }}_t:={\prod }_{s=0}^t{\alpha }_s$。根据公式$(1)$可以得到$t$时刻的噪化状态$x_t$与$t-1$时刻的噪化状态$x_{t-1}$的关系为：
$$x_t=\sqrt{{\alpha }_t}\cdot x_{t-1}+\sqrt{1-{\alpha }_t}\cdot {ϵ}_{t-1},\text{\hspace{1em}(4)}$$其中：${ϵ}_{t-1}\sim \mathcal{N}(\text{0},\text{I})$，通过不断取代递推可以得到$t$时刻的噪化状态$x_t$与输入$x_0$之间的关系为：
KaTeX parse error: No such environment: equation* at position 8: \begin{̲e̲q̲u̲a̲t̲i̲o̲n̲*̲}̲ \begin{aligned…其中：$ϵ\sim \mathcal{N}(\text{0},\text{I})$， ${\overline{ϵ}}_{t-2}$是两个高斯分布相加后的分布。第一步到第二步的公式推导需要说明一下，根据高斯噪声的特点，对于两个方差不同的高斯分布$\mathcal{N}(\text{0},{\sigma }_1^2\cdot \text{I})$和$\mathcal{N}(\text{0},{\sigma }_2^2\cdot \text{I})$，其相加后的高斯分布为$\mathcal{N}(\text{0},({\sigma }_1^2+{\sigma }_2^2)\cdot \text{I})$，表现在公式中，即：
KaTeX parse error: No such environment: equation at position 8: \begin{̲e̲q̲u̲a̲t̲i̲o̲n̲}̲ \begin{aligned…其中：两个高斯分布相加后的标准差为：
$$\sqrt{{\alpha }_t(1-{\alpha }_{t-1})+(1-{\alpha }_t)}=\sqrt{1-{\alpha }_t{\alpha }_{t-1}},\text{\hspace{1em}(7)}$$依此得到第二步，进而逐渐递推到最后一步。公式$(3)$和公式$(5)$的目的就是表明在前向扩散过程中，由于每步加的噪声均是同分布的高斯噪声，因此不需要逐步进行加噪，直接就可以由输入$x_0$的到$T$时刻的噪化状态$x_T$。当${\overline{\alpha }}_T\approx 0$，$T$时刻的分布$x_t$则几乎就是一个高斯分布，据此其可以定义为：
$$q(x_T):=\int q(x_T|x_0)q(x_0)d{x}_0\approx \mathcal{N}(x_T;\text{0},\text{I}),\text{\hspace{1em}(5)}$$其中：$\int$表示积分，最终的噪化状态$x_T$也可以在图像上看出其分布特点。

3. 反向采样

反向采样过程就是根据已有的噪化状态通过学习来估计噪声分布，进一步获得上一时刻的状态，并逐渐从高斯分布中构造出真实数据。根据前向扩散过程的结果，可以认为$T$时刻噪化状态$x_T$的后验分布$p(x_t)\sim \mathcal{N}(x_t;\text{0},\text{I})$，则联和分布$p_{\theta }(x_0,x_1,...,x_T)$也是一个马尔科夫链，其被定义为:
$$p_{\theta }(x_0,x_1,...,x_T):=p(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t),\text{\hspace{1em}(6)}$$则$t-1$时刻的噪状态$x_{t-1}$可以由上一时刻$t$的状态$x_t$得到，其条件分布可以表示为：
$$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\sum }_{\theta }(x_t,t)),\text{\hspace{1em}(7)}$$其中：${\mu }_{\theta }(x_t,t)$和${\sum }_{\theta }(x_t,t))$分别为$t$时刻由噪声估计网络得到的噪声均值和方差，$\theta$为噪声估计网络的参数。此时，在输入为$x_0$时，$t-1$时刻的噪状态$x_{t-1}$与上一时刻$t$的状态$x_t$之间的真实条件分布为：
$$q(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};{\tilde{\mu }}_t(x_t,x_0),{\tilde{\beta }}_t\cdot \text{I}),\text{\hspace{1em}(8)}$$其中：噪声后验分布参数${\tilde{\mu }}_t$, ${\tilde{\beta }}_t$分别为：
$${\tilde{\mu }}_t=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}\cdot {ϵ}_t),{\tilde{\beta }}_t=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}\cdot {\beta }_t,\text{\hspace{1em}(9)}$$此处认为${\sum }_{\theta }(x_t,t)={\sigma }_t^2\cdot \text{I}$，即${\sigma }_t^2={\tilde{\beta }}_t$，则预测的后验条件分布变为：
$$p_{\theta }(x_{t-01}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\sigma }_t^2\cdot \text{I}),,\text{\hspace{1em}(10)}$$即利用噪声估计网络${\mu }_{\theta }$来估计真实噪声分布均值${\tilde{\mu }}_t$，则公式$(9)$中的噪声分布均值可以被估计为：
$${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}\cdot {ϵ}_{\theta }(x_t,t)),\text{\hspace{1em}(11)}$$而根据公式已知$t$时刻的噪化状态$x_t$满足$x_t=\sqrt{{\overline{\alpha }}_t}\cdot x_0+\sqrt{1-{\overline{\alpha }}_t}\cdot ϵ$，则网络学习的优化目标就是让估计出的噪声分布接近真实的噪声分布，即：
$${\mathbb{E}}_{x_0,t,{ϵ}_t\sim \mathcal{N}(0,\text{I})}[||{ϵ}_t-{ϵ}_{\theta }({\sqrt{\overline{\alpha }}}_t\cdot x_0+\sqrt{1-{\overline{\alpha }}_t}\cdot ϵ,t)|{|}^2],\text{\hspace{1em}(12)}$$而$t-1$时刻的噪化状态$x_{t-1}$可以表示为 (这块尚没搞清楚这个公式的由来，似乎与原论文中的公式不一样)：
$$x_{t-1}=\sqrt{{\overline{\alpha }}_{t-1}}(\frac{x_t-\sqrt{1-{\overline{\alpha }}_t}\cdot {ϵ}_{\theta }(x_t,t)}{\sqrt{{\overline{\alpha }}_t}})+\sqrt{1-{\overline{\alpha }}_{t-1}}\cdot {ϵ}_{\theta }(x_t,t),\text{\hspace{1em}(13)}$$其中：$z\sim \mathcal{N}(\text{0},\text{I})$。则根据不同时刻噪声估计网络对噪声分布的估计可以依据公式$(13)$逐渐反向采样得到真实数据分布。

3. 图像条件扩散模型

在图像恢复任务中，必须使用条件扩散模型才能生成我们预期的恢复图像，实际中即将退化的图像作为条件引入到噪声估计网络中来估计条件噪声分布。如图所示：

图像条件扩散模型与经典扩散模型的前向扩散过程完全一样，区别仅在于反向采样过程中是否引入图像条件。则反向采样过程中$x_1,x_2,...,x_T$的联合分布变为：
$$p_{\theta }(x_0,x_1,...,x_T|\hat{x}):=p(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t,\hat{x}),\text{\hspace{1em}(14)}$$其中，$\hat{x}$为作为条件输入噪声估计网络的退化图像。此时，噪声分布估计变为：
$${ϵ}_{\theta }(x_t,t)\to {ϵ}_{\theta }(x_t,\hat{x},t)，\text{\hspace{1em}(15)}$$$t-1$时刻的噪化状态$x_{t-1}$也由公式$(13)$变为：
$$x_{t-1}=\sqrt{{\overline{\alpha }}_{t-1}}(\frac{x_t-\sqrt{1-{\overline{\alpha }}_t}\cdot {ϵ}_{\theta }(x_t,\hat{x},t)}{\sqrt{{\overline{\alpha }}_t}})+\sqrt{1-{\overline{\alpha }}_{t-1}}\cdot {ϵ}_{\theta }(x_t,\hat{x},t),\text{\hspace{1em}(16)}$$
实际中，条件的引入由多种方式，最常见的方法是直接通过与噪化状态拼接后作为噪声估计网络的输入。

4. 可以考虑改进的点

以下是我问chatGPT得到的答案：

我的拙见：

1. 引入天气退化图像恢复中：虽然扩散模型已经出现众多研究，但在图像去雨、去雾、去雨滴、去雪等方面的研究屈指可数；
2. 改进噪声估计网络：经典的扩散模型是基于U-Net结构的，主要模块也是卷积 (也包括自注意力)，近来有一些研究发现Transformer架构在扩散模型上可以取得更好地效果；
3. Follow最新的更快地扩散模型：传统的扩散模型要进行图像恢复，一幅图片的处理时长基本都是几十秒，实时性太差，目前有一些研究提出快速反向采样的方法；
4. 无监督：目前多数给予扩散模型的图像恢复算法仍然是有监督的 (当然不算是监督学习，只是条件生成)，可以采用一些无监督策略来利用扩散模型实现图像恢复。

5. 实现代码

参考的论文是Restoring Vision in Adverse Weather Conditions with Patch-Based Denoising Diffusion Models，有能力调通源代码的建议查看原文的github链接。我的代码是在其基础上进行了大量简化，将没有必要的代码删减，方便理解和查看。简化后的用于图像恢复的扩散模型代码见：完整可直接运行代码，其中包括详细的实验操作流程，只需要修改数据集路径即可直接使用。