RainDiffusion:When Unsupervised Learning Meets Diffusion Models for Real-world Image Deraining

RainDiffusion:When Unsupervised Learning Meets Diffusion Models for Real-world Image Deraining (Paper reading)

Mingqiang Wei, Nanjing University of Aeronautics and Astronautics, China, arXiv, Cited:0, Code: 无, Paper.

1. 前言

当无监督学习遇到真实世界图像去雨的扩散模型时会发生什么？为了回答这个问题，提出了RainDiffusion，这是第一个基于扩散模型的无监督图像去雨范式。除了传统的无监督的图像去噪智慧之外，RainDiffusion引入了对未配对的真实世界数据的稳定训练，而不是弱对抗性训练。RainDiffusion由两个合作分支组成：非扩散翻译分支（NTB）和扩散翻译分支。NTB利用循环一致性架构，通过生成初始干净/下雨图像对来绕过标准扩散模型的不成对训练的困难。DTB利用两个条件扩散模块来逐步细化具有初始图像对和扩散生成先验的期望输出，以获得更好的去噪和降雨生成的泛化能力。RainDiffusion是一种非对抗性训练范式，是真实世界图像去噪的新标准。大量实验证实了我们的RainDiffusion方法优于非/半监督方法，并显示了其优于完全监督方法的竞争优势

2. 整体思想

提出了一种新的基于去噪扩散概率模型（DDPM）的无监督学习范式，称为RainDiffusion，以解决现实世界中普遍存在的图像去噪问题。整体思想是用两个扩散模型生成无噪和有噪的图像，然后将有噪的图像作为条件concat到无噪扩散模型的去噪器，以获得无噪图像。

3. 方法

图2:RainDiffusion的pipeline。RainDiffusion由两个主要的交互分支组成：非扩散翻译分支（NTB）和扩散翻译分支。详细地说，仅在模型训练阶段使用的NTB充分利用循环一致性架构来生成初始的干净/下雨图像对。请注意，$G_{\Phi }^A$和$G_{\Phi }^B$分别用于生成雨天和无雨图像。$x,y,x^{{}^{\prime }},y^{{}^{\prime }},x^{{}^{\prime \prime }},y^{{}^{\prime \prime }}$是指给定的无雨图像、给定的雨天图像、生成的雨天图像，生成的无雨图、重构的无雨像和重构的雨天图像。给定初始图像对，DTB利用两个条件扩散模块在干净图像和雨天图像之间进行转换。

3.1 Non-diffusive Translation Branch (NTB)

尽管条件扩散模型已经显示出最先进的图像条件数据合成和恢复能力，但扩散模型的实现仍然需要成对的干净/下雨图像。我们将通过设计一个新的非扩散翻译分支（NTB）来解决这个问题，而不需要对抗训练，NTB引入了循环一致性架构来生成用于条件扩散模型训练的初始干净/下雨图像对。由于训练的雨天图像$y$和干净图像$x$是不成对的，NTB使用了两个参数为${\Phi }^{A,B}$的非扩散生成器以获得初步翻译估算：

这里不使用判别器，两个生成器是不是很像自编码器？？损失函数为：$L_{cyc}=E_{x\sim P_{data(x)}}[||x^{{}^{\prime \prime }}-x|{|}_1]+E_{y\sim P_{data(y)}}[||y^{{}^{\prime \prime }}-y|{|}_1]$
，有了这两个网络，我们就可以获得图像对了。

3.2 Diffusive Translation Branch (DTB)

继去噪扩散概率模型在生成建模领域的优势之后，DTB设计了两个条件扩散模块，以减轻不稳定训练的限制，提高真实世界图像去噪和降雨生成的泛化能力。

训练过程采用的是条件逆扩散过程，既把条件concat到去噪器中。因为自然图像都是尺寸不一的，因此采用基于Patch的扩散模型。Patch大小选择128。为了加速训练和采样，使用DDIM加速算法。算法流程如下：首先输入的是清晰图像和有雨图像，使用$Crop$将图像分为Patch，这里的$G$对应着两个生成器，也就是这两个扩散模型（这里的训练过程一定是需要采样的，一定特别慢，但是他这里没说明！！）。损失函数第一行：将生成的有雨图像作为条件，使用扩散模型A去雨；损失函数第四行，将重构的有雨图像作为条件，使用扩散模型A去生成的雨；损失函数第二行，将重构的无雨图像作为条件，使用扩散模型B生成雨图；损失函数第三行，将生成的无雨图像作为条件，使用扩散模型B生成真实雨图；损失函数最后一行为有监督学习的约束。

训练完成后，只使用扩散模型A来去雨，而对应的有雨图像对是扩散模型B获得的。采样过程如下：$D$对应着单张图片分块数，${\hat{\Omega }}_t$指这些重叠的噪声估计累积在矩阵中它们各自的贴片位置。$M$基于每个像素的接收估计数对${\hat{\Omega }}_t$进行归一化，并使用平滑的全图像噪声估计${\hat{\Omega }}_t$执行隐式采样更新。

4 实验

使用Nvidia GeForce RTX 3090 GPU的系统上使用Pytorch 1.6实现RainDiffusion。为了优化RainDiffusion，我们使用最小批量为4的Adam优化器来训练范式，其中动量参数${\beta }_1$和${\beta }_1$的值分别为0.5和0.999。此外，生成器和鉴别器的初始学习率都设置为0.00002。在实验中，我们将${\lambda }_{cyc}$设置为1:0。对于训练，从原始图像（或其水平翻转版本）中随机裁剪128×128图像块。

5. 总结

本文想法与SR3类似，只是本文处理的是真实世界的图像，这导致了复杂度急剧上升。本文的方法太过臃肿复杂。效果跟SR3相差也不是很多。