基于离散注意力机制的信息蒸馏去雨网络研究

罗 岩，周 杰*，罗 宏，徐 蕾，邵根富

(1.南京信息工程大学 电子与信息工程学院，江苏 南京 210044；2.杭州电子科技大学 自动化学院，浙江 杭州 310000)

0 引言

近年来，由于深度学习的快速发展，单幅图像训练和其他低水平视觉任务取得显著进展.单图像去雨旨在从被雨水污染的图像中恢复干净的图像，这是其他下游计算机视觉任务的基础，例如目标检测、图像分类、人员识别等.尽管人们对图像去雨任务的关注和改进越来越多，但由于灾难性遗忘问题，现有的基于深度卷积网络的雨水去除方法只能在单一类型的数据集上学习成对的雨水/清洁图像之间的固定映射规则.此外，通过对多个数据集进行顺序训练，当前的深度神经网络几乎完全忘以前的知识，并且在对新任务进行训练后，在很大程度上降低模型在以前任务中的性能.

本文旨在解决图像去雨问题.去雨是一个高度不适定的问题，为解决此问题，许多关于背景图像或雨层的先验知识[1]被提出.虽然这些方法在一定程度上解决去雨任务，但由于对降雨的形状、方向和大小的简单假设，当先验无效时，它们总是无法工作.得益于更强的特征表示能力，卷积神经网络在许多计算机视觉问题上取得巨大成功，例如，目标检测、目标跟踪、姿势估计、图像修复[2]、语义分割、去模糊、去雾[3]，超分辨率以及去雨.这些去雨算法考虑不同的去雨方式.研究人员设计各种深度网络结构来解决图像去雨问题.为提高计算效率和减少模型尺寸，轻量级网络相继问世.通过学习雨水模型或估计雨水密度来指导去雨过程，雨水检测和清除[4]和多头网络[5]被提出，并呈现出更好的去雨性能.其他各种网络模块是通过设计各种机制开发的，例如，通过挤压和激发的通道注意[6]，使用非局部平均网络的像素级注意[7]和对抗性学习[8].尽管这些去雨方法取得更好的性能，但它们仍然忽略一些重要的细节.

以往的工作主要集中在特征提取和处理或神经网络结构上，而目前的雨水去除方法已经可以取得显著的效果，基于单一网络结构的训练而不考虑交叉尺度关系可能会导致信息丢失.尽管这些传统方法在某些理想情况下被证明是有效的，但它们主要受到主观的、人工设计的先验知识和巨大的计算负担的限制.近年来，由于深度神经网络(Deep Neural Networks，DNN)强大的非线性拟合能力，基于深度学习的方法有助于视频去雨任务的显著改进.该方法的核心思想是，基于合成的数据集，以端到端的方式直接训练由DNN参数化的去雨模型.这些方法大多利用不同的技术，例如超级像素对齐[9]、双水平流[10]和自学习[11]，从雨天视频中提取干净的背景.尽管这些基于深度学习的方法在一些综合基准上取得了令人印象深刻的去雨效果，但仍有很大的空间来进一步提高性能和实际应用中的泛化能力.

为解决这个问题，本文提出一种基于离散注意力机制的信息蒸馏去雨网络，用于从单个图像中去除雨水.首先，构建一个多尺度特征提取模块，通过建立各个尺度之间的相关性来融合不同尺度的特征，以便更好地学习雨水特征.其次，为使网络增加对不同尺度的雨纹的学习能力，引入循环网络.最后通过一种跨通道融合网络来学习不同尺度下的特征，这种方式将不同通道上的特征连接起来，充分利用通道注意力机制，将不同通道的特征融合，提升模型在不同降雨场景的鲁棒性.

1 相关工作

1.1 图像特征点提取

对于实例级识别任务来说，尺度变化是一个众所周知的挑战.在经典的计算机视觉应用以及在基于深度学习的方法中，构建金字塔表示一直是跨各种图像分辨率处理视觉信息的有效方法.

由于特征提取的运算复杂度和可复用度都不太相同，到目前为止图像特征点仍缺乏万能或者准确的描述.特性的准确描述通常根据需求和应用类型确定.不同大小和形状的目标在进行相同的降采样比例后通常会产生很大的语义代沟问题，最普遍的体现就是小目标检测准确度比较低.

特征金字塔在不同尺寸下具有不同分辨率的特征，各种尺寸的目标都能够在适当的尺寸下具有正确的特征表示，并且通过融合更多尺寸信息，在不同尺寸下对各种尺寸的目标都做出估计，从而很好地改善模型的特性.多尺度深度卷积神经网络(Multi-scale Deep Convolutional Neural Network，MSCNN)[12]利用多级特征映射进行预测，但没有跨不同的特征级别进行聚合.特征金字塔网络(Feature Pyramid networks，FPN)[13]是当前通过自上而下路径和横向连接学习不同层次特征表示的领先范式.类似地，反向连接的对象先验网络(Reverse Connection with Objectness Prior Networks，RON)[14]引入了反向连接，将信息从高级功能传递到低级功能.

本文通过多尺度空间卷积池化模块跨通道提取特征，融合不同尺度的特征，丰富提取的图像信息.

1.2 注意力机制

近几年来，注意力机制在深度学习的方方面面都被证明有对模型优化的巨大进步.许多方法开始关注注意机制.Wang等[15]提出了一种剩余注意网络，该网络具有使用下采样和上采样的空间注意机制.此外，挤压与激励网络(Squeeze-and-Excitation Networks，SENet)[6]提出了通道注意机制.它在通道上执行，然后使用完全连接层计算每个通道的权重.其中，卷积注意力模块(Convolutional Block Attention Module，CBAM)[16]声称，由于信息丢失，全局平均池化(global average pooling，GAP)只能获得次优特征.为了解决这个问题，它使用了GAP和全局最大池，并获得了显著的性能改进.同样，基于空域富模型的隐写分析方法(Spatial Rich Model，SRM)[17]还建议将GAP与全球标准偏差池结合使用.在CBAM的推动下，全局二阶池化卷积网络(Global Second-Order Pooling Convolutional Networks，GSoP)[18]引入了一种二阶池法进行下采样.

在神经网络的学习过程中，模型拥有的参数量越大，模型获得的信息量越大，精细化表达能力越强，但参数过多有可能导致信息过载.为解决这一问题，引入注意力机制，通过增加对当前任务更重要的参数的权重，降低其他无关参数的相关性来提高模型的处理精确度和速度.因此本文提出离散注意力信息蒸馏网络使网络能精细化处理不同雨纹的不同特征，提升去雨效果.下面详细介绍基于离散注意力机制的信息蒸馏去雨网络(Multi-Split Distillation Attention Network，MSDAN)的整体结构以及损失函数和模型算法.

2 本文方法

2.1 MSDAN整体框架

本文提出的基于单幅图像的去雨网络采用多尺度空间特征提取、多重蒸馏离散注意力结构，使卷积神经网络在高的空间分辨率下能够提取更加精细的雨痕特征.通过将包含大规模信息的低分辨率分支提取的特征传入高分辨率分支中，用于引导雨痕特征提取.

在多重蒸馏离散注意力网络中，采用特征图分割块，上下文聚合增强块，多重信息蒸馏模块.特征图分割块用于增强网络对特征图处理的效率，上下文聚合增强块用于提高去雨图像在不同分辨率下的细节信息，在去雨性能和计算机资源消耗之间寻求更好的平衡，多重信息蒸馏模块通过提取部分细特征，并将剩余的粗特征放入下一层进行深度处理来增强收集的各种精细信息.设计基于结构相似性(Structural Similarity，SSIM)损失函数，利用Adam优化方法提高网络去雨效果.MSDAN的结构如图1所示.

图1 MSDAN网络架构

2.2 多空间卷积池化网络

通常来说在提取图像特征时，既想让特征具有较大的感受野，又想让分辨率保持不变，以免丢失图像边缘的信息.一个比较常用的处理方式就是图像金字塔，先将原图调整到各个尺度，再输入到相同的网络，得到不同的特征图，再做融合.这个方式确实能够提高精确度，然而模型处理速度会大幅降低.而空洞卷积在获取最大感受野的同时，也可使分辨率损失减小.

为解决这一问题，本文设计多尺度空间卷积池化网络模块(Multi-Split Network，MSN)，见图2所示.通过采用具有不同速率的空洞卷积进行特征提取，将不同尺度的信息融合在一起，提高不同图像的细节信息.为此将不同空洞卷积层的得到的特征进行单独处理，整合融合特征.

图2 多空间卷积池化网络架构

2.3 离散注意力信息蒸馏网络

离散注意力信息蒸馏网络(Dispersed Attention Info-Distillation Network，DAID)通过将通道方向的注意力归纳为特征地图组表示，见图3所示.具体是构建出一个由特征图组和分裂注意力操作组成的模块(DA).则先将特征图分割，分别通过R个卷积核处理，最后将对应通道的结果乘以权重来获得最终结果，其中R为每个输入的通道数.与此同时，把他们产生的特征图整合起来(特征融合).H、W、C为块输出特征图大小.

图3 DAID网络架构

这里第c个分量的计算如下：

(1)

第c通道的计算公式为：

(2)

(3)

式(3)中：sk为第c通道后每次分割的权重.

与以前大多数使用递归结构的小参数模型不同，本文构建出一个信息蒸馏网络(Info-Distillation，ID)，见图4所示.

图4 信息蒸馏模块

引入深度可分离卷积的思想，通过提取多个蒸馏步骤的特征来对之前的特征进行分裂处理，使网络轻量化，组成网络的信息蒸馏模块.同时通过提取细节特征，并将浅层特征和深层特征融合，使网络在实现轻量化的同时，增加重要特征的权重.为聚合所有步骤提取的特征，加入一个对比度感知通道注意层，特别是与低级视觉任务相关的，以增强收集的各种精细信息.

具体来说，利用更有用的特征(边缘、角点、纹理等)进行图像恢复.保留的部分可以被视为细化的特征.考虑到Fin的输入特性，第n个信息蒸馏模块中的这个过程可以描述为：

(4)

下一阶段是将每个步骤中的细化功能连接起来.它可以表示为：

(5)

式(5)中：Concat表示沿通道维度的连接操作.

2.4 循环网络

为加强不同尺度间雨纹特征的关联，并且增强MSDAN对不同特征的学习效率，选择引入循环卷积网络来加强不同特征间的联系，优化网络去除不同雨纹的能力.在众多循环网络中，选择(Long Short Term Memory networks)LSTM结构进行特征提取，因为在图像去雨方面，LSTM比循环神经网络(RNN)，门控循环网络(GRU)更具有经验优势.

2.5 损失函数

许多损失函数都被用于训练单幅图像去雨网络，如MS-SSIM等损失函数，有些时候超参数调优会变得很困难.本文引入负SSIM损失函数训练网络.具体可以表示为：

L=-SSIM(XT,Xgt)

(6)

式(6)中：XT代表第T个循环的结果，其中Xgt为对应的基准干净图像.

(7)

3 实验结果与分析

3.1 实验环境及评价指标

本文算法软件环境为Python3.7和深度学习框架Pytorch,硬件环境为Tesla A40处理器，48 G内存.开发工具为PyCharm 专业版2021.3.3.图像尺寸大小为100×100，批处理大小设置为8，采用Adam优化器，实验迭代次数为100次，初始学习率设为1×10-3.每到30、50和80个迭代次数时将学习率乘以0.2进行衰减.

本文采用两个指标作为评估算法性能的参照，分别为峰值信噪比(PSNR)和结构相似性(SSIM)，去除雨纹后的结果还原度越高，两个度量指标的值就越大.

3.2 数据集

当前数据集大部分的雨线分布单一，去雨模型训练的时候很容易对相同的噪声产生拟合，为测试网络的泛化性，针对上述问题，本文构建合成数据集RS1000.合成数据集共1 000张雨图，通过从Rain200L和Rain200H数据集中随机提取1 000张雨图作为真实数据集，图像的分辨率为321×481.并从其中随机提取900张雨图做训练集，剩余100张做训练集，并将其取名为RS1000.

现实生活中，有雨和无雨的图像很难大量获得，因此本文选用两个常用数据集和一个自建数据集，分别为Rain100L，Rain100H，RS1000.Rain100L数据集有200张雨图做训练集，100张雨图做测试集.Rain100H有1 800张雨图作为训练集，其中Ren 等[19]为防止网络过拟合，将Rain100H 中相似的图像删除至1 254张图，由200张雨图做测试集.

3.3 对比实验

在SSIM和PSNR两个指标下，将提出的方法与LPNET[20]、RESCAN[21]、PReNet[19]和DCSFN[22]进行比较.采用3个合成数据集测试模型，分别为Rain100H和Rain100L和RS1000.并将得到的定量结果与相应数据集下的测试方法进行比较.所有的对比结果都是严格按照原文作者提供的代码和参数进行实验得出的，保证客观性和公平性.对比的去雨方法的比较结果如表1所示.

表1 去雨效果对比

由表1可以清楚地看到，所提出的方法在所有数据集中实现最高的SSIM和PSNR.在三个数据集上，MSDAN比其他去雨网络在PSNR指标上分别高出0.60～7.16，1.13～5.3，1.32～11.44；在SSIM指标上分别高出0.013～0.091，0.004～0.024，0.011～0.085.结果表明，相比RESCAN用膨胀卷积获得感受野，MSDAN拥有更强的夸大感受野的能力，相比于PReNet使用特征融合极大的丰富特征信息，相比于DCSFN拥有更强的细化特征提取能力.图5展示各个算法去雨的视觉效果.LPNet去雨后仍存在大量阴影，而RESCAN、PReNet、DCSFN处理后的雨图虽然相对清晰，但仍然有不少白斑.这四种模型在恢复效果上都存在一定的不足，而MSDAN既能很好的去除雨纹又能较为满意的恢复纹理细节，这进一步证明所提方法的优越性.

图5 不同方法的可视化对比结果

3.4 消融实验

MSDAN引入阶段性递归，即递归残差网络，以显著减少渐近网络的参数，递归次数为T.对MSDAN的递归次数T在不同的情况下优化效果进行测试，见表2所示.所用的测试数据为Rain100L.递归次数分别为T=3,4,5,6,7，在递归次数 T=6时，PSNR和SSIM的值都取得最好且优于其他递归次数.选取MSDAN3、MSDAN4、MSDAN5、MSDAN6、MSDAN7的输出图像，见图6所示.结果显示，递归次数不是越深越好.当递归次数小于6时去雨效果会变差.证明MSDAN递归次数为6时最优.

表2 比较不同递归次数的MSDAN模型

图6 不同递归次数的可视化对比结果

为探究各模块对模型总体去雨的效果影响，对模型各模块进行消融测试.首先，使用6个连续的递归残差网络作为基础的去雨网络.通过在每个残差块中添加信息蒸馏网络(ID)，使网络更好的细化提取到的雨纹和背景特征.由于雨纹的方向、密度不同，本文在每个信息蒸馏模块前都添加离散注意力模块(DA)来重点提取雨纹信息.通过在DAID前添加多尺度空间卷积池化网络(MSN),使网络更好的细化提取图像特征点.同时为避免网络过深会造成提取特征丢失，在MSN后添加长短时间记忆网络(LSTM)加强不同特征的联系，使用Rain100L公开数据集进行测试.实验分别测试LSTM和MSN的存在时，网络的优化效果见表3所示.

表3 不同组合模块融合学习

模型通过基线网络进行对比，依次添加DA模块、ID模块、MSN模块、LSTM模块进行消融实验，展现各个模块在图像去雨中的优势，证明了本文算法提高图像去雨的可行性.

4 结论

本文提出一种基于离散注意力机制的信息蒸馏去雨网络MSDAN来解决图像去雨问题.通过多尺度特征融合方法进行特征提取，使用信息蒸馏网络来精准提取特征.同时利用离散通道注意力捕捉重要特征，以实现不同尺度下信息的内部连接和充分利用.根据Rain100H、Rain100L和RS1000数据集上的实验结果显示，相比其他几种深度学习方法，本文模型在PSNR和SSIM等评价指标上优于其他算法，但仍存在一定不足.算法中缺少对空间注意力的利用，同时时间成本也需要优化，这也将是下一步研究的方向.