这篇文章提出了一种名为DRSformer的稀疏Transformer网络,用于解决图像去雨问题。主要内容总结如下:
-
问题背景:
图像去雨是一个典型的低级视觉问题,旨在从雨天图像中恢复出清晰的图像。现有的Transformer方法虽然能够建模非局部信息,但在处理局部细节时表现不佳,且标准自注意力机制可能会引入不相关的特征干扰。 -
主要贡献:
-
稀疏Transformer架构:提出了一种稀疏Transformer网络(DRSformer),通过自适应地保留最有用的自注意力值来改善特征聚合,从而提升去雨效果。
-
Top-k稀疏注意力(TKSA):设计了一种可学习的top-k选择算子,动态保留每个查询中最关键的注意力分数,避免不相关信息的干扰。
-
混合尺度前馈网络(MSFN):开发了一种多尺度前馈网络,探索多尺度特征以更好地去除雨纹。
-
专家混合特征补偿器(MEFC):引入了专家混合特征补偿器,通过多种稀疏CNN操作提供协作精炼,进一步提升去雨性能。
-
-
实验与结果:
在多个公开数据集(如Rain200L/H、DID-Data、DDN-Data和SPA-Data)上进行了实验,结果表明DRSformer在PSNR和SSIM指标上优于现有的最先进方法,尤其是在处理复杂雨纹时表现出色。 -
消融研究:
通过消融实验验证了top-k选择、MSFN和MEFC的有效性,证明了每个组件对最终去雨性能的贡献。 -
局限性:
尽管DRSformer在去雨性能上表现出色,但其模型复杂度较高,未来计划通过剪枝或蒸馏技术进行模型压缩。
DRSformer通过稀疏注意力机制和多尺度特征融合,显著提升了图像去雨的效果,并在多个基准数据集上取得了优异的性能。这里是自己的论文阅读记录,感兴趣的话可以参考一下,如果需要阅读原文的话可以看这里,如下所示:
官方项目地址在这里,如下所示:
摘要:
基于Transformer的方法在图像去雨任务中取得了显著的性能,因为它们能够建模非局部信息,这对于高质量图像重建至关重要。本文发现,大多数现有的Transformer通常使用来自查询-键对的所有相似度进行特征聚合。然而,如果查询中的标记与键中的标记不同,从这些标记估计的自注意力值也会参与特征聚合,从而干扰清晰图像的重建。为了解决这个问题,我们提出了一种有效的去雨网络——稀疏Transformer(DRSformer),它能够自适应地保留最有用的自注意力值进行特征聚合,从而使聚合的特征更好地促进高质量图像重建。具体来说,我们开发了一种可学习的top-k选择算子,以自适应地保留每个查询中最关键的注意力分数,从而实现更好的特征聚合。同时,由于Transformer中的简单前馈网络没有建模多尺度信息,而这对于潜在清晰图像的重建非常重要,我们开发了一种有效的混合尺度前馈网络,以生成更好的特征用于图像去雨。为了学习一组丰富的混合特征,结合了来自CNN算子的局部上下文,我们为模型配备了专家混合特征补偿器,以提出一种协作精炼去雨方案。在常用基准上的大量实验结果表明,所提出的方法在性能上优于现有的最先进方法。
1. 引言
单幅图像去雨是过去十年中出现的一个典型的低级视觉问题。它旨在从观察到的雨天图像中恢复出干净的图像。由于清晰图像和雨纹是未知的,这是一个不适定的逆问题。为了解决这个问题,早期的方法通常基于雨纹和清晰图像的统计特性施加各种先验。事实上,这些手工设计的先验对于复杂和多变的雨天场景并不鲁棒,这限制了去雨性能。
最近,许多基于学习的方法[4, 19, 23, 36, 52, 53, 56]已经将各种CNN架构作为与传统算法相比的优选选择。然而,卷积操作的内在特性,即局部感受野和输入内容的独立性,阻碍了模型消除长程雨退化扰动的能力。为了缓解这些限制,Transformer[2, 26, 35, 50]已被应用于图像去雨,并取得了不错的性能,因为它们能够更好地建模非局部信息以实现高质量图像重建。然而,这些方法在恢复清晰图像时,图像的细节(即图像的局部特征)并没有得到很好的建模,如图1所示。一个主要原因是Transformer中的自注意力没有建模局部不变特性,而CNN在这方面做得很好。
2. 相关工作
单幅图像去雨:由于图像去雨是一个不适定问题,传统方法[12, 20, 24, 30, 60]通常开发各种图像先验以提供额外的约束。然而,这些手工设计的先验往往依赖于经验观察,因此无法建模清晰图像的固有特性。为了解决这个问题,许多基于CNN的框架[53]被开发出来用于解决图像去雨问题,并取得了不错的恢复性能。为了更好地表示雨分布,一些研究考虑了雨的特性,如雨的方向[27]、密度[61]、雾化效应[15],并通过递归计算[19, 23, 36]或转移机制[16, 49, 54, 55]优化网络结构。尽管这些方法比基于手工先验的方法取得了更好的性能,但由于卷积的内在限制,它们难以捕捉长程依赖关系。与基于CNN的去雨方法不同,我们利用Transformer作为网络骨干来建模非局部信息以实现图像去雨。
视觉Transformer:受Transformer[7]在自然语言处理(NLP)[40]和高级视觉任务[1, 28]中的巨大成功启发,Transformer已被应用于图像恢复[2, 13, 48, 51, 58],并且比之前的基于CNN的基线表现更好,因为它们能够建模非局部信息。在图像去雨领域,Jiang等人[18]设计了一种动态关联去雨网络,通过将Transformer中的自注意力与背景恢复网络结合。最近,Xiao等人[50]精心开发了图像去雨Transformer(IDT),通过基于窗口和空间的双重Transformer实现了出色的结果。需要注意的是,大多数现有方法依赖于密集的点积自注意力作为Transformer的核心。然而,这种计算方式的一个缺点是,具有较小权重的冗余或不相关特征可能会干扰注意力图,这使得输出特征包含潜在的噪声。在这项工作中,我们提出了Transformer中的稀疏注意力,以缓解标准自注意力所面临的最相关信息被忽略的问题。
稀疏表示:受生物大脑中神经活动的启发,深度神经网络中隐藏表示的稀疏性作为一种诱人的“免费午餐”在视觉和NLP任务中崭露头角[44, 64]。事实上,稀疏表示在处理低级视觉问题(如图像去雨[46]和超分辨率[31])中也起着关键作用。原则上,稀疏注意力可以分为基于数据(固定)的稀疏注意力和基于内容的稀疏注意力[6, 38]。对于基于数据的稀疏注意力,几种局部注意力操作被引入到CNN骨干中,主要考虑仅关注局部窗口大小。最近的研究[11, 42]已经研究了在Transformer骨干中强制执行稀疏性。最近,Zhang等人[63]设计了一种注意力可伸缩Transformer,允许来自稀疏区域的标记交互特征,这是基于数据的稀疏性。与之不同,我们基于top-k选择实现了一种简单但有效的自注意力近似,以实现稀疏注意力,这是基于内容的稀疏性。
Top-k选择:Zhao等人[64]首次在NLP任务中提出了一种基于top-k机制的显式选择方法。受其成功的推动,k-NN注意力[42, 44]被进一步引入以提升视觉Transformer。与在空间维度上执行top-k选择[44]不同,我们设计了一种高效的top-k有用通道选择算子。
图2. 所提出的用于图像去雨的稀疏Transformer网络(DRSformer)的整体架构,主要包括以下部分:
-
稀疏Transformer块(STB):包含top-k稀疏注意力(TKSA)和混合尺度前馈网络(MSFN)。
-
专家混合特征补偿器(MEFC):用于提供协作精炼。
其中:
-
LN 表示层归一化(Layer Normalization)。
-
DW-Conv 表示深度卷积(Depth-wise Convolution)。
3. 提出的方法
在本节中,我们首先描述了图像去雨的整体流程和对称分层网络架构。随后,我们提供了所提出的稀疏Transformer块(STB)的详细信息,作为我们方法的基本构建单元,它主要包含两个关键元素:top-k稀疏注意力(TKSA)和混合尺度前馈网络(MSFN)。最后,我们介绍了引入的专家混合特征补偿器(MEFC)。
稀疏Transformer块:由于标准Transformer[7, 40, 58]使用所有标记全局计算自注意力,这对于图像恢复不友好,因为它可能涉及不相关特征之间的噪声交互。为了解决这些限制,我们开发了一个稀疏Transformer块(STB)作为特征提取单元,利用神经网络中出现的稀疏性[64]。形式上,给定第(l-1)块的输入特征Xl−1,STB的编码过程可以定义为:
4. 实验与分析
实验设置:
数据集:我们在多个公共基准上实现了去雨实验,包括Rain200L/H [52]、DID-Data [61]和DDN-Data [8]。Rain200L和Rain200H包含1,800张合成的雨天图像用于训练和200张用于测试。DID-Data和DDN-Data由12,000和12,600张具有不同雨方向和密度级别的合成图像组成。分别有1,200和1,400张雨天图像用于测试。此外,我们还使用了一个大规模的真实世界数据集SPA-Data [45]来评估我们的方法,该数据集包含638,492对图像用于训练和1,000对用于测试。
比较方法:我们将DRSformer与两种基于先验的模型(DSC [30]和GMM [24])、基于CNN的方法(DDN [8]、RESCAN [23]、PReNet [36]、MSPFN [19]、RCDNet [43]、MPRNet [59]、DualGCN [9]和SPDNet [56])以及最近的基于Transformer的方法(Uformer [48]、Restormer [58]和IDT [50])进行了比较。对于最近的代表性方法(DualGCN、SPDNet、Restormer和IDT),如果未提供预训练模型,我们使用作者提供的模型重新训练,否则我们使用其在线代码进行评估以确保公平比较。对于其他方法,我们参考了[10, 50]中报告的一些结果。
评估指标:我们采用PSNR [34]和SSIM [47]作为上述基准的评估指标。根据之前的去雨方法[10, 19],我们在YCbCr空间的Y通道中计算PSNR和SSIM指标。对于没有真实图像的雨天图像,我们使用无参考指标,包括NIQE [33]和BRISQUE [32]。
与最先进方法的比较:
合成数据集:表1报告了在不同基准数据集上的定量评估。如图所示,我们可以注意到我们提出的方法在所有去雨方法中表现最好,特别是在PSNR方面,例如,DRSformer在平均上比并发方法IDT高出0.4 dB。与之前的基于CNN的模型相比,这一进展更加明显。在DID-Data和DDN-Data基准上的显著增加分数表明,我们的方法能够正确处理各种类型的空间变化雨纹。为了提供令人信服的证据,我们在图3中展示了由最近方法生成的样本的视觉质量比较。纯基于CNN的模型,例如MPRNet和SPDNet,在重雨场景中无法恢复清晰图像。可以看到,所有基于Transformer的计算方法在细节和纹理恢复方面都存在缺陷。不幸的是,IDT甚至引入了相当大的边界伪影。得益于开发的具有top-k选择的稀疏注意力,我们的方法可以生成高质量的去雨结果,这些结果与真实图像更加一致。
真实世界数据集:我们进一步在SPA-Data基准数据集上进行了实验,相应的结果在表1的最后一列中提供。正如预期的那样,我们的模型继续实现了最高的PSNR/SSIM值,展示了DRSformer在去雨性能和泛化方面的优越性。视觉质量比较可以在图4中观察到。相比之下,我们的方法在去除大多数雨纹的同时保留了真实的图像结构方面显著优于其他方法。为了进一步验证DRSformer的有效性,我们还从Internet-Data [45]中随机选择了20张没有真实图像的真实雨天图像进行另一次评估。如表2所示,我们的网络获得了较低的NIQE和BRISQUE值,这意味着在真实雨天场景中,与其他比较模型相比,我们的网络输出了更清晰的内容和更好的感知质量。通过图5中的定性质量比较,大多数深度学习模型对空间长雨纹敏感,并留下一些明显的雨效应。相反,我们的网络成功地去除了大多数雨扰动,并具有视觉上令人愉悦的恢复效果,这意味着它可以很好地泛化到未见过的真实世界数据类型。
消融研究:
为了理解这种top-k选择的效果,我们进一步使用高通滤波(HPF)来可视化学习到的特征,如图8所示。与标准自注意力操作(无top-k)相比,我们的策略可以更好地帮助重建更精细的特征并提高潜在的恢复质量。由于附近的像素往往比其他像素更相似,top-k选择算子有助于减少来自长程像素依赖的不相关上下文。这一选择步骤允许在自注意力计算过程中丢弃较小相似度权重(来自部分长程特征交互),从而促进更准确的表示以实现高质量输出。
MSFN的有效性:为了评估提出的MSFN的有效性,我们将其与三个基线进行了比较:(1)常规前馈网络(FN)[7],(2)Doony前馈网络(DFN)[25],(3)门控深度卷积前馈网络(GDFN)[58]。在Rain200H上的定量分析结果列于表3中。尽管GDFN在两个相同尺度的深度卷积流中引入了门控机制以带来性能优势,但它仍然忽略了去雨的多尺度知识。通过在不同尺度上添加局部特征提取和融合,MSFN确实可以更好地提升性能,并在GDFN的基础上实现了0.21 dB的PSNR增益。
MEFC的有效性:为了评估MEFC的有效性,我们基于不同的模型变体进行了实验,如表4所示。与基线模型(a)相比,MEFC提供了额外的性能优势,这要归功于辅助数据稀疏性。此外,我们观察到,网络管道中不同位置的MEFC对恢复性能有特定影响。事实上,我们还分析了每个MEFC中不同专家数量的影响。当使用单一专家模型(d)时,与我们的多专家模型(f)相比,性能显著下降。与将所有专家设置为相同结构[21]不同,我们的多专家公式更加多样化,这带来了性能增益,因为不同的感受野和不同的CNN操作。通过图9中的放大框,具有上述所有组件的模型的恢复结果往往更清晰,因为它能够在恢复过程中充分利用更多样化的特征。总的来说,我们的模型(f)比其他可能的配置表现更好,这表明我们考虑的每个设计策略都对DRSformer的最终性能有其贡献。
密切相关的方法:我们注意到,最近的方法[22]提出了一种k-NN图像Transformer(KiT)来解决图像恢复问题,通过聚合k个相似补丁与成对局部注意力。与KiT相比,KiT采用了复杂的局部敏感哈希,无法确保足够的全局交互,而我们简单但有效的top-k选择机制不仅享受局部性,还增强了全局关系挖掘的能力。由于KiT的代码不可用,我们参考了其论文的结果。图10显示了在Rain800 [62]上训练的定性比较。我们可以看到,KiT往往会模糊内容并导致颜色失真。相比之下,我们的方法导致了更好的去雨结果。
此外,我们还注意到[44]最近设计了k-NN注意力,通过选择top-k相似标记来增强视觉Transformer的表示能力。与KVT [44]不同,KVT在空间维度上实现top-k选择,我们的算子在跨通道计算稀疏注意力方面更高效。此外,我们提出的TKSA中的稀疏度kk是动态可学习的,而不是[44]中的固定设置。在这里,我们采用KVT中的kk-NN注意力来替换我们的TKSA进行比较。为了确保公平比较,模型测试保持了相同的训练设置。如图10(c)和(d)所示,我们的方法可以生成更清晰的图像。
5. 结论
我们提出了一种有效的稀疏Transformer网络,称为DRSformer,以解决图像去雨问题。基于观察到Transformer中的标准自注意力可能会受到不相关信息的全局交互的影响,我们开发了top-k稀疏注意力以保留最有用的自注意力值,以实现更好的特征聚合。为了促进去除雨纹的聚合特征,我们开发了一种混合尺度前馈网络,以更好地探索多尺度表示。此外,专家混合特征补偿器被引入到模型中,以为稀疏Transformer骨干提供协作精炼,从而保留重建图像的精细细节。实验结果表明,我们的DRSformer在性能上优于最先进的方法。
局限性:我们提出的方法旨在进一步提高图像去雨性能,但在模型效率方面存在局限性。具体来说,我们的模型需要33.7百万个参数,并在大小为256×256的图像上消耗242.9G FLOPs。我们将在模型中应用剪枝或蒸馏方案,以在保持原始去雨性能的同时实现可信的模型压缩。