SDNET: MULTI-BRANCH FOR SINGLE IMAGE DERAINING
USING SWIN
前言
最近，流行的transformer具有全局计算特性，可以进一步促进图像去雨任务的发展。本文首次将Swim-transformer引入图像去雨领域，研究了Swim-transformer在图像去雨领域的性能和潜力。具体来说，我们对Swim-transformer的基本模块进行了改进，设计了一个三分支模型来实现单幅图像的去雨。前者实现了基本的雨型特征提取，而后者融合不同的特征进一步提取和处理图像特征。此外，我们还采用jump connection来融合深层特征和浅层特征。实验表明，现有的公共数据集存在图像重复和背景相对均匀的问题。因此，我们提出了一个新的数据集Rain3000来验证我们的模型。
Transformer[28]最初是自然语言处理（NLP）领域的一个模型，用于并行处理单词向量，以加速模型推理。它的全局计算特性适用于远距离传递特征。这正是计算机视觉领域中卷积运算所不擅长的。Dosovitskiy等人[29]将图像分割成16x16个图像块，将不同的图像块作为不同的词输入到transformer中，提高了图像分类的精度。近年来，人们从深度[30]、多尺度[31]等角度应用transformer来完成相关任务。然而，Transformer也有不可忽视的缺点，例如计算量与图像大小之间存在二次关系，这限制了它的应用环境。Liu等人[32]提出的Swin-transformer使用滑动窗口使模型具有线性计算复杂度，通过跨窗口连接改善了窗口间的信息交换，最终提高了模型在图像分类、目标检测和实例分割等方面的性能。
本文提出了一种新的图像去雨网络SDNet，它是利用Swim-transformer强大的特征表示能力构建的端到端去雨网络。具体地说，我们改进了Swim-transformer的基本模块，重新设计了一个双分支模型，实现了单图像去雨。前者实现了基本的雨型特征提取，后者融合了不同分支的特征。此外，我们采用jump connection来融合深度特征和浅层特征，以提高网络模型的性能。
本文贡献如下：
- 改进了Swim-transformer模块，使基本块的输出大小与特征提取的输入一致。
- 为了提高特征融合的效率，提出了一种三分支网络模型，并利用额外的jump connection来融合深特征和浅特征，以提高网络的学习能力。
- 提出了一个全新的数据集。并利用该数据集对本文提出的模型和其他主流降雨模型进行了训练和检验。
- 实验表明，与主流去雨模型相比，该模型具有一定的性能和推理速度优势
最近有大量的研究工作将transformer引入CV域，并取得了良好的效果。具体来说，Dosovitskiy等人[29]将图像分成16X16个图像块，然后将其拉伸成一维向量，然后送入网络中完成图像分类任务。Chen等人[38]提出了一种基于卷积运算的transformer与Unet相结合的TransUnet方法，实现医学图像的分割。蒋等[39]设计了与对抗生成网络结构相同的图像生成transformer。transformer中的self-attention导致模型计算直线增长，导致transformer不能在低计算能力的硬件上运行。Liu[32]提出了一种利用滑动窗口方法使网络计算线性增长并加速网络推理的方法。我们的方法是基于这种方法来实现一个单一的图像去雨任务的融合特征。
本文的方法是基于这种方法[32]来实现一个单一的图像去雨任务，融合不同分支的特征、深度特征和浅层特征。
方法
Swin-transformer
Transformer是一个功能强大的网络模块，可以取代CNN操作。但其中的Muti-Head Attention导致模型的计算量迅速增加，导致transformer模型无法在许多底层硬件中测试和使用，注意力的数学表达式如下：
其中 是query, key and value matrices。 表示number of patches in a single Windows。d表示dimension of query or key。B中的值取自偏差矩阵 . Swin-transformer提出通过固定窗口大小来限制单个自注意计算所涉及的参数量，然后通过滑动窗口来实现窗口间的信息交换。Swin-transformer块如图1所示。Swin-transformer使用特征压缩来获得高层语义特征，并在图像分类、目标检测和图像分割方面取得最佳效果。损失且无法恢复，因此此模型不能直接用于从单个图像中去雨（这句话没懂...）。
Network structure
本文使用一个简单而强大的前馈网络作为主干，如图2所示。SDnet网络基本上由三个多分支融合模块组成，称为MSwt，一个多分支模块MSwt-m和两个基本block模块。此外，还增加了跳转连接，目的是融合深特征和浅特征，以提高网络去雨的性能。
受Swin-transformer的启发，本文的Basic-block模块设计如图3所示。输入图像或特征图的大小为231× 231× 3，并且在PatchPartition中patch block的大小设置为3× 3，因此矩阵的size of the position encoding 是 。这是因为在图像去雨任务中更多的关注像素信息和位置信息，而在图像分类等高级任务中更多的关注语义信息，所以VIT[29]和文献[40]的patch blocks都被设置为16× 16和location encoding的大小为 .在in Linear Embedding 中，将拉长的向量映射为其长度的两倍，并进行full concatenation，以避免特征压缩。在Swin-transformer块中，窗口大小被设置为7，heads of self-attention数被设置为3。我们在Swin-transformer块之后添加了一个Unpatch Partition operation，这是一个reciprocal operation with Patch Partition ，使得输出图像或特征的大小和尺寸与输入图像或特征的大小和尺寸相同。在Swin-transformer块中，W-MSA和SW-MSA共同完成全局特征提取。本文中，Basic-block使用一个集合。
为了更灵活地构建网络，提出了Basic-block的概念，并设计了两个三分支特征融合块。如图4和图5所示，与后者相比，前者有一个用于融合特征的附加基本块。数学表达式如下：
其中F（·）表示基本块的操作。x表示模块Mswt的输入。这种设计的思想来源于自我注意中的多头注意机制。通过学习F1、F2、F3，可以自适应地学习不同的特征。将输入映射到不同的子空间，分别提取不同的特征。与自我注意不同的是，我们对提取的特征求和，而不是级联操作。通过F4融合增加的特征，实现进一步的特征提取。由于设计思想来源于多头注意机制，多分支具有与该机制相同的特点，即在一定范围内，分支越多，模型性能越好。为了平衡模型的规模和模型的性能，我们选择了三个分支进行特征提取。
虽然transformer可以保持特征在长距离传播，但是仍然需要在网络中结合深特征和浅特征，为此我们设计了一个没有特征融合的Mswt模块，我们称之为Mswt-m，如图5所示，其数学表达式如下：
F1、F2、F3将输入映射到三个不同的空间进行特征提取，对提取的特征求和，然后与第二个Mswt模块的输出求和，再经过一个基本块，实现深度特征和浅层特征的融合，如图2中的小跳跃连接所示，而图2中相对较长的跳跃连接则考虑了主要特征中包含的丰富的空间和纹理信息，有助于完成深度特征中缺失的纹理信息。
Loss function
其中，O为雨图像，B为对应标签。是绝对差（SAD）之和，用于计算相似预测图像和标签之间的像素损失，如等式6所示。SSIM（结构相似性）是结构相似性，最初用作评估两个图像内容的结构相似性的度量。Ren等人[41]证明了SSIM作为损失函数在图像降额任务中的有效性的负面作用，其数学表达式如等式7所示。尽管使用该损失函数可以获得高SSIM度量，但图像仍然存在失真和低峰值信噪比（PSNR）。identity loss（等式8）由CycleGAN[42]导出，CycleGAN[42]用于约束生成图像的颜色丢失，这里我们使用它来约束图像去雨后的图像样式，这减少了图像失真，提高了网络性能。α , β , λ 是SAD损失、SSIM损失和identity loss的系数。在本文中，分别设置为0.2、4和1。
Experimental environment
实验使用Tesla V100 16G GPU进行训练，使用Pytorch框架1.7.0和（Adam）[43]，初始学习率为5× 10−4，减少到5× 10−5和5× 10−6当训练迭代次数分别为总迭代次数的3/5和4/5时。输入模型的图像大小设置为231×231. batch size为5。
我们提出了一个全新的数据集用于网络训练和消融实验。该数据集是从ImageNet中随机抽取的10万幅图像，保证了图像的多样性。从Efficientderain[12]降雨模式数据集中随机选择一到四种降雨模式，并添加到选定的图像中。我们最终选择了3000张合成图像作为训练集，400张作为测试集。我们把这个数据集命名为Rain3000。此外，我们还使用公开的数据集Rain100L和Rain100H[44]来验证SDnet模型。两个公开的数据集都包含1800个训练图像和200个测试图像。
使用SSIM和PSNR作为评价指标，这两种指标已被广泛用于评价预测图像的质量。PSNR是根据两幅图像之间的像素误差来计算的，误差越小，值越大，图像越相似，除雨效果越好。相反，图像去雨的效果越差
结论
首先，本文提出了一种基于Swin-transformer的三分支端到端除雨网络，它充分利用了Swin-transformer强大的学习能力，用一种改进的Swin-transformer代替卷积运算，并设计了一个多分支模块来融合不同空间域的信息，使用跳转连接来融合深特征和浅特征。此外，我们提出了一个新的数据集，由3000个训练对和400个测试对组成。该数据集是基于ImageNet生成的，具有丰富的背景和雨型组合，便于模型的推广。我们提出的模型在数据集Rain3000和公共数据集Rain100L、Rain100H上都达到了最佳性能。我们的工作还有些不足。例如，在参数数目相同的情况下，哪种方法更适合于并行或串行的图像去噪任务还没有详细探讨。以及是否可以使用多个不同大小的滑动窗口来实现窗口间的进一步信息交换，以提高网络降容的性能。此外，我们正在使用更简单的前馈网络，更复杂的网络仍然值得研究