DRT: A Lightweight Single Image Deraining Recursive Transformer

摘要

过度参数化是深度学习中常见的技术，以帮助模型学习和充分概括给定的任务；然而，这往往导致巨大的网络结构，并在训练中消耗大量的计算资源。最近在视觉任务上强大的基于Transformer的深度学习模型通常有很重的参数，并承担着训练的难度。然而，许多密集预测的低级计算机视觉任务，如去除雨痕，在实践中往往需要在计算能力和内存有限的设备上执行。因此，我们引入了一个基于递归局部窗口的自注意结构，并提出了去雨递归Transformer（DRT），它具有Transformer的优越性，但需要少量的计算资源。特别是，通过递归结构，我们提出的模型在去雨中只使用了目前表现最好的模型的1.3%的参数数量，同时在Rain100L基准上超过最先进的方法至少0.33dB。消融研究还调查了递归对去雨结果的影响。此外，由于该模型不是刻意为去雨设计的，它也可以应用于其他图像复原任务。我们的实验表明，它可以在去雪上取得有竞争力的结果。

1、介绍

计算机视觉是一个快速增长的领域，在现实世界中有许多应用，如用于安全的面部识别、用于自主车辆的物体检测和用于字幕生成的场景理解。这些高层次的计算机视觉任务通常需要在未被破坏的图像上进行操作以获得最佳性能。在雨天拍摄的图像自然会被破坏，因为雨中的条纹会阻挡背景，导致信息的丢失。因此，找到一种有效的方法来去除图像中的雨痕并重建背景，同时保持模型的简单性，对于高层视觉任务在雨天的操作至关重要，特别是对于计算资源有限的设备。

                       

在过去的五年中，许多模型都广泛采用卷积神经网络（CNN）来进行单幅图像去雨。特别地，Fu等人设计了一个基本的轻量级CNN，将图像的高频部分作为输入，以去除雨；所得到的图像被添加回图像的低频部分，很像一个剩余连接。在这项工作之后，采用了更先进的CNN架构以达到更好的效果。Fan等人使用多个残差连接和递归CNN块来进行雨痕去除。Li等人利用非局部增强的密集块，并表明具有精心设计的结构的深度网络可以增强训练的鲁棒性并取得更好的结果。Yang等人很好地总结了2020年之前的单幅图像去除方法，我们参考该文献调查以了解详细情况。最近，许多视觉任务已经由基于注意力机制的架构很好地完成了；然而，它们可能在计算上很繁重，训练起来也很耗时，正如在[19]中讨论的那样。 然而，它们的计算量很大，训练起来也很费时，这一点在下一节将详细讨论。这限制了视觉Transformer可以应用的领域的范围。

为了给图像重建任务设计一个高效的深度视觉Transformer，我们提出了一个带有跳过连接的递归注意力结构，称为Deraining递归Transformer（DRT）。特别是，一个精心设计的递归结构在网络中的某些自注意力块之间共享权重，因此在增加网络深度时，参数的数量不再急剧增加。此外，在图像去雨中，输入图像和输出图像有很大一部分重叠的信息；因此，通过利用从输入到输出的剩余连接，只需要少量的自我注意块就可以集中检测雨痕，从而减少网络的大小。有了这些设计，我们的模型只需要1.18M个参数（即目前单幅图像脱隐技术的1.3%），就能在Rain100L基准上比其他脱隐方法至少多出0.33dB。图1中显示了不同方法的去伪存真结果的例子，这表明所提出的DRT有很强的能力来重建各种不规则的背景。此外，与[6]提出的ViT或[3]提出的IPT不同，它们需要在数百万张图像上进行训练，而DRT只需要在一个小数据集（700张图像）上轻松训练，因此有效降低了训练难度。综上所述，这项工作有以下贡献：

- 据我们所知，这是研究递归结构视觉Transformer用于图像重建任务的第一项工作。
- 我们的模型在单幅图像重建任务上的表现优于其他许多方法，同时消耗的计算资源较少。
- 所提出的模型具有足够的通用性，可以执行其他的图像修复任务，如单幅图像解密。

2、相关工作

2.1、视觉Transformers

自从Vaswani等人引入Transformers网络后，人们做了很多研究，用Transformers取代CNN作为视觉任务的骨干结构。特别是Dosovitskiy等人引入了视觉Transformers（ViT），将图像分割成不重叠的小斑块，并将每个斑块作为原始Transformers的一个标记。由于图像补丁的嵌入机制，ViT在分类任务上表现良好，同时保持了相对较低的计算复杂性。为了扩展ViT的能力，Liu等人提出了Swin-Transformer。Swin Transformers设计背后的目的之一与我们工作的目的相同，即降低视觉变换器的计算复杂性。这种类型的Transformers在局部窗口中执行多头的自我注意，而不是像ViT中那样全局的。计算成本从减少到，其中图像维度为，局部窗口大小用表示。

同时，Swin-Transformer采用了移位窗口法来补偿全局信息的损失。一个Swin Transformer Block（STB）首先将输入通过一个层规范（LN），然后通过一个带窗的多头自我注意（WMSA）机制。自我注意机制的输出与初始输入相结合，形成第一条剩余连接。然后，输出通过一个多层感知器（MLP），再通过另一个LN。最后的输出再次与自我注意机制的输出相结合，形成第二个残差连接。多个STB被用来处理一个给定的图像，并在识别、检测和分割任务中取得最先进的结果。

2.2、将视觉Transformers应用到图像复原上

上一节的视觉Transformer及其变体最近被利用来实现更好的一般图像修复结果。Chen等人[3]使用基于[19]的Transformer编码器-解码器结构来设计一个被称为图像处理Transformer（IPT）的通用图像修复变换器。同样地，Wang等人提出使用一个Transformer块，即LeWin块，以及一个带有跳过连接的沙漏结构来进行背景重建。这种设计在不同的图像重建任务中表现良好，同时保持相对较低的计算成本。此外，Liang等人在Swin-Transformer的基础上设计了残差转化器块，并将多个残差块堆叠起来，为图像重建进行深度特征提取。与IPT不同的是，IPT是在ImageNet上预训练的，这种类型的Transformer是用800个样本训练的，同时相对于许多其他基于CNN或基于Transformer的方法，保持了少量的网络参数。

 2.3、递归网络的结构

深度学习中的递归结构会重复调用网络中的某部分层来处理数据。递归结构在图像重建任务中的应用可以追溯到[10]。在这项工作中，我们提出了一个深度递归的CNN，以稳定训练过程并保持一个参数很少的深度结构。它的经验表明，一个递归层就足以实现良好的图像超分辨率结果。后来，同样的想法被[7]用来设计一个残差引导的递归CNN网络，用于图像去伪存真。我们的方法与这项工作相似，因为我们也彻底研究了在基于Transformer的架构上使用递归和残差结构，而不是依靠CNN。虽然递归结构也会导致潜在的问题，如爆炸性梯度或过拟合，但精心设计的递归结构可以导致最佳的网络结构，这一点将在以下章节中讨论。

3、方法

3.1、网络结构

DRT由三个阶段组成：块嵌入阶段、深度特征提取阶段和图像重建阶段。 图2给出了它的可视化。基本上，我们可以把第一阶段看作是输入预处理阶段，第三阶段是推理输出后处理阶段，而第二阶段是主要的推理阶段。给定一个原始的雨季图像，下面的方程就抓住了DRT的功能：

                                   

块嵌入阶段首先对雨季图像应用卷积层，并将其映射到，将图像分成补丁并按深度叠加。这里，分别为原始图像的高度、宽度和深度。表示用于将输入分成不重叠的小窗口的补丁尺寸，是嵌入维度，通常是。注意//号表示整数除法操作。由于Transformer通常是用一维token工作的，我们重新塑造卷积层的输出来完成网络的第一阶段。然后，这个输出可以被送入转化器块进行处理。我们可以简明扼要地将第一阶段写为：

                                   

 

下一阶段由多个递归转换块（RTB）组成，共同进行方程3所述的深度特征提取。我们用来表示RTB的总数。RTB的细节将在下一小节给出。这里我们注意到，这个阶段保持着输入维度，每个RTB之间没有参数共享。这个阶段的输出被表示为。在这个过程的最后，一个残差连接被用来将深度特征提取阶段的输入添加到其输出中，以限制RTBs专注于检测雨水条纹。

                   

 深层特征，，然后由图像修复阶段进行处理，如公式4所示，它颠覆了第一阶段的过程。输入图像的高度和宽度首先通过重塑来恢复，卷积层用来处理这些特征，得到输出。最后，在网络的输入和输出之间使用了另一个剩余连接，这样网络提取的雨痕特征就可以从输入中去除。

                     

 3.2、递归Transformer块

每个RTB使用递归和剩余连接来堆叠Swin Transformer Blocks（STB）。在数学上，STB可以描述为：

                                 

 在这里，WMSA可以用公式[19]来概括：

                             

                             

在上述公式中，指的是权重。输入的、和，通常被称为查询、键和值。术语指的是K的维度。如果我们把每个剩余连接中的STB数量表示为，递归调用的数量表示为，让输入表示为，那么每个RTB可以用以下等式描述。

                                              

符号用于表示多个函数组成。输入和输出的尺寸保持不变。剩余连接总是从RTB的输入开始。上式中的第一个控制RTB的递归调用，第二个控制不同STB的组合。剩余连接中的每个可以有不同的参数，但递归调用总是调用相同的STB，如上式所示，因为索引从未被使用。例如，在图2中，每个剩余连接中有两个STB（），这种连接在一个RTB中被递归调用三次（）。在每个RTB的末尾，卷积层被用来处理基于局部窗口的Swin Transformer可能忽略的信息。如果采用单一卷积层，我们在卷积后不使用任何激活。对于多个卷积层，我们使用Leaky ReLU作为激活函数。

3.3、损失函数

有许多提出的损失函数来促进网络的训练；然而，我们发现使用平均平方误差损失是足够的。对于一个输出图像和其对应的地面真实图像，误差由以下公式给出：

                            

其中表示条目总数，表示第个条目的地面真实像素值，表示第个条目的预测值。对于多个预测，我们只需对每一对的误差进行平均，即可得到最终的误差测量。

 4、实验

我们通过在三个不同的数据集上比较DRT与其他先进技术的减损结果（PSNR和SSIM）以及所需的参数数量和计算成本来测试DRT。然后，我们进行了一项消融研究，以研究架构的变体。最后，我们表明DRT可以执行其他的图像修复任务，如图像去噪。

4.1、数据集

在以下数据集上进行了实验，以评估拟议的DRT的有效性。

Rain800

这个数据集包括700张训练图像和100张测试图像（测试100）。这些图像选自UCID数据集和BSD500训练集，并在上面合成了雨的条纹。不同的雨水条纹包含在Rain800中。

Rain100L

在这个数据集中有1800张训练图像和100张测试图像。背景图像选自BSD200数据集。每张图像都由一个方向的浅色雨水条纹组成。

Rain100H

与Rain100L类似，该数据集由1800张图像和100张测试图像组成。背景也选自BSD200。然而，每张图像可以由多个方向的大雨条纹组成，提高了减损性能。

Snow 100K

背景图像从Flickr API中下载，有三种类型的雪的密度（即小、中、大）随机选择并合成在背景图像上。训练集由5万张图片组成。三个测试集（即Snow100K-S、Snow100K-M、Snow100K-L）对应于三种雪的密度。Snow100K-S包含16611张图像，Snow100K-M包含16588张图像，Snow100K-L有16801张图像。

4.2、设置和训练

设置
表现最好的模型具有图2所示的网络结构。补丁嵌入和图像重建阶段仅由一个没有任何激活的卷积层组成。在深度特征提取阶段有六个RTB（N=6），每个RTB由两个STB上的三个递归调用（L=3）组成。在每个RTBs的末尾只采用了一个单一的卷积，没有任何激活函数。所有的卷积操作都保持了输入的大小，因此在去粗取精的过程中，没有对图像进行缩小或提高采样，这可能会导致像素级信息的丢失。对于每个STB，局部窗口的尺寸被固定为77，补丁大小被设定为1，头数为2。我们的分层特征表示的深度被选为D=96。

训练
最初的训练过程主要是在Rain-800数据集上进行的，该数据集包括700个具有不同类型雨痕的合成雨天训练样本和100个用于网络验证的测试数据（test100）。我们没有从头开始训练，而是在其他两个数据集上对这个网络进行微调。训练过程对消融研究来说略有不同，在下面的章节中讨论。训练时使用了一块8GB内存的RTX 2070超级显卡。我们使用随机裁剪，裁剪尺寸为5656，并对训练集进行随机水平翻转，以进行数据扩充。我们为所有的训练过程选择了Adam优化器。批量大小为8，初始学习率被锁定为1e-4。在微调阶段，学习率被调整到1e-5和1e-6之间的数值。网络被训练到误差在100个历时中没有下降，或者在最后50个历时中误差只减少了1e-3。我们得到的最好的模型被训练了4600个epochs，在这个模型上进一步训练会导致过量。由于许多其他基于Transformer 的网络在训练过程中消耗了大量的数据，我们还尝试了一个更广泛的数据集，即Rain 13k，这是一个不同的derain数据集的混合物，共包含13712个雨天-清洁的图像对。然而，我们没有看到网络的性能有任何明显的变化。一个小的数据集对这个模型来说是足够的。                   

       

4.3、定量结果

表1显示了在三个数据集上评估的最佳表现的DRT模型的PSNRs和SSIMs，以及与其他先进技术的比较结果。我们的结果以粗体字显示，以便更好地进行可视化。对于Rain100L，在PSNR指标下，我们的性能超过了表1中列出的所有方法，至少有0.33db。对于同一指标，我们的方法在Test100和Rain100H数据集上的表现排名第二。同样，在SSIM指标下，DRT的结果在Rain100L上排名第二。另一方面，从表1中可以看出，我们的方法在所使用的参数数量方面是非常轻便的。具体来说，与HiNet相比，DRT只使用了1.3%的参数数量。请注意，早期的方法没有发布源代码，我们在表中省略了他们的参数。一般来说，在参数数量和脱因性能之间存在一个权衡。请参考图5，了解这种权衡的视觉效果。在这个图中，为了更好的视觉效果，具有大量参数的模型被忽略了。靠近左上角的方法倾向于平衡参数数量和脱隐结果，反之亦然。DRT通过最接近左上角的方法很好地平衡了这种权衡，而Jorder-E和HiNet等方法则以大的网络规模为代价获得良好的性能。其他方法可能使用较少的参数，但不能达到与DRT一样好的脱隐结果。我们进一步为HiNet、PreNet和DRT提供了一个固定的输入，以计算它们相应的乘数-累加器运算（MAC）（HiNet也使用这个指标）以及存储网络和执行前向和后向传递所需的内存量。结果显示在表2中。DRT消耗的MAC最少，同时享有比HiNet大约少4.5倍的内存使用。

                      

                     

4.4、定量结果

我们选择在Test100、Rain100L和Rain100H上评估的不同方法的视觉结果，并分别在图3、图4和图1中展示。选择这些图像是为了显示不同雨量类型的导出性能。在图3中，雨的密度很高，但每条雨痕都非常小。DRT能够去除这些雨痕中的大部分，而Hi-Net则忽略了其中的许多雨痕。请注意，在输入的背景和地面实况（GT）图像的背景中存在着色调的变化。由于训练数据没有考虑背景颜色的变化，我们修复的图像不能恢复原始色调。在图4中，其他方法似乎将一些雨点条纹解释为土地或云层的一部分，因此不能完全去除这些雨点，而DRT可以去除几乎所有的雨点。在图1中，我们可以看到由火车产生的烟雾中存在伪影，而由Jorder和Hi-Net产生的图像的烟雾中也存在伪影。我们的模型和PreNet对这些复杂的软边形式给出了更好的重建。最后，我们还呈现了一个真实的derain结果，如图6所示。与Hi-Net相比，DRT可以去除这些雨点中的大部分，从放大的图片中可以看出。请注意，对于现实的图片，由于训练数据没有考虑雾霾和雨水的混合，所以由雨水引起的雾霾仍然是预期的。

4.5、超参数调试

我们发现，最关键的设计选择是减少STBs使用的patch大小。原始的Swin Transformer 块将补丁大小设置为4。当将补丁大小减少到1时，我们的PSNR有了显著的增加。这表明，在图像提取中考虑本地像素信息的重要性。此外，我们还发现，在整个网络中保持输入维度（即不对输入进行任何上采样或下采样）是有益的。由于单幅图像的衍生基本上是一个密集的预测任务，这两个设计选择确保了网络推断的像素级信息没有损失。该模型的其他超参数也得到了广泛的调整，我们在此对它们进行总结。如前所述，我们没有发现对原始Swin Transformer中的移窗方法的使用有任何影响；这可能是由于我们在每个RTB的末端使用了卷积，并对每个STB进行了局部位置编码，以考虑局部和全局的位置信息，而移窗方法是为了解决原始Transformer中全局信息的损失。同样地，我们尝试将变换器的工作维度深度从96增加到180，但在性能上没有改善。我们试图放弃输入和输出之间的剩余连接，但没有观察到任何改善。这可能是因为残余连接限制了推理阶段对雨痕的检测，而雨痕在输入图像中只占据了一小部分信息。

4.6、消融研究

然后，我们对DRT的基本设计特征进行了消融研究，即RTB的数量（N）、STB的数量和STB上的递归调用数量（L）。在这项研究中，每个模型都是用类似的训练程序在Rain100L上从头训练。所有不包括在消融研究中的其他超参数都被固定在最佳网络结构上。每个模型的评估是在Rain100L数据集上进行的。表3中列出了每个评估的参数数量和PSNR，最佳模型显示在最后一行。从表中可以看出，前两行研究了RTB数量对最佳DRT模型的影响；然而，相应的低PSNR结果表明，增加和减少RTB的数量都极大地限制了DRT的性能。第三行到第五行研究STB的递归调用数量对DRT性能的影响。再次，无论是增加还是减少调用次数，都不能提高DRT的脱隐性能。请注意，只有一个递归调用是指没有任何递归的DRT的变体。然而，它的性能是所有其他变体中最差的，这表明递归在设计中的重要性。然后我们同时调整RTB的数量和递归调用的数量，但结果并不理想。因此，我们得出结论，有6个RTB和3个递归调用对于有希望的脱隐性能是必不可少的。与此相反，表3的第二行和倒数第三行表明，STB数量的变化对脱嵌结果影响不大。值得注意的是，只有一个STB可以将参数的数量从1.18M降低到一百万以下，即0.841M，同时仍然保持良好的脱嵌结果。

4.7、单图像去雪上的结果

我们进一步测试了DRT在单幅图像解密任务上的表现，以经验性地展示其解决不同图像修复任务的能力。最好的模型是在Snow100K数据集上训练的，训练过程与之前相同。然后在Snow100K提供的相应的三个数据集上测试该模型；每个数据集包括一种雪的类型。不同方法的PSNR结果显示在表4中。我们模型的性能在这些方法中排名第一或第二。我们还在图7中展示了一个desnow视觉结果的样本，以及相应的输入、输出和遮罩图像。虽然输入图像中的雪很密集，但DRT可以去除大部分的雪，并恢复一个非常接近地面实况的干净背景。雨痕和雪花在物理上和视觉上都有很大的不同，这意味着它们需要不同的功能来去除；这项研究表明，经过训练，DRT并不局限于单一图像的去除，也可以很好地应用于其他领域。

5、结论

基于Transformer的网络通常计算量很大。我们利用基于局部窗口的自我注意机制、残差和递归连接，以轻量级的架构实现图像去污的雨痕。它在Rain100L数据集上的表现最好，在其他数据集上也保持了一流的性能。由于其简单的结构，它使用了少量的参数，同时对计算资源的依赖性较低。由于该模型中没有刻意设计去除雨痕，我们表明它可以应用于除雪，并有可能应用于其他图像重建任务。