论文来源：ICLR2019

论文链接：pdf (openreview.net)

项目链接：GitHub - yulunzhang/RNAN: PyTorch code for our ICLR 2019 paper "Residual Non-local Attention Networks for Image Restoration"

作者：美国东北大学

概述：

提出了一种残差非局部注意力网络，可以在图像去噪，去马赛克，压缩伪影去除，超分辨率任务上取得较好的效果。

网络整体采用模块堆叠的方式，由2个卷积层，2个残差非局部注意块（RNAB）和8个残差局部注意块（RAB）组成，并以残差学习的方法重建图像。其中，RNAB和RAB的结构相似，以RAB为例，每个RAB开头和结尾都有2个残差块（RB），中间是一个残差注意力的结构。RB的结构为：conv-relu-conv-res，是一个简化版的残差块，中间的残差注意力模块有两个分支组成，主干分支和掩模分支，主干分支同样由2个RB块组成，掩模分支结构为：RB-conv(stride>1)-RB-RB-反卷积-RB-conv-sigmoid，主干分支和掩模分支提取到的特征相乘再和之前的特征相加作为结尾2个RB的输入。如果是RNAB，则需要在RAB的开头位置加一个非局部块（NLB），NLB的输入为一个H*W*2C的特征图，先由三个分支的1*1卷积对特征图维度做变换，使得他们能做矩阵乘法和加法，先用前两个分支得到的矩阵做一次矩阵相乘，再用sigmoid函数来得到每个像素位置与所有像素位置的相似性，进而可以得到一个全局相似性矩阵，再用这个矩阵与第三个分支的矩阵相乘来融合全局信息，最后把得到的矩阵再通过一个1*1的卷积来进行reshape后与输入相加。

ABSTRACT

本文提出了一种用于高质量图像恢复的残差非局部注意力网络。由于没有考虑图像中信息的不均匀分布，以往的方法受到局部卷积运算以及空间或通道特征同等处理的限制。为了解决这个问题，我们设计了局部和非局部注意块来提取特征，这些特征能够捕捉像素之间的长依赖关系，并且更加关注具有挑战性的部分。具体来说，我们在每个（非）局部注意块中设计了主干分支和（非）局部掩码分支。主干分支用于提取层次特征。局部和非局部掩模分支的目的是在混合注意的情况下自适应地重新缩放这些层次特征。局部掩模分支关注更多的卷积运算局部结构，而非局部注意力更多地考虑整个特征图中的长程依赖关系。此外，本文还提出了残差局部和非局部注意学习来训练超深网络，进一步增强了网络的表征能力。我们提出的方法可以推广到各种图像恢复应用，如图像去噪，去马赛克，压缩伪影减少，超分辨率。实验结果表明，该方法在定量和直观上均优于目前的主流方法。

提出了一种用于高质量图像恢复的超深残差非局部注意力网络（RNAN）。我们设计了残差的局部和非局部注意块作为超深网络的基本构建模块。每个注意块由主干和掩模分支组成。我们介绍了残差块用于主干分支和提取层次特征。对于掩模分支，我们采用大步长卷积和反卷积的方法进行特征降尺度和升尺度，以扩大感受野的大小。此外，我们在掩模分支中加入非局部块以获得剩余的非局部混合注意力。我们将RNAN应用于各种恢复任务，包括图像去噪、去噪和压缩伪影减少。大量的实验表明，我们提出的RNAN在所有的任务中都达到了目前领先的结果。据我们所知，这是第一次考虑剩余非局部注意的图像恢复问题。

这项工作的主要贡献有三个方面：

• 我们提出了非常深的残差非局域网络用于高质量的图像恢复。强大的网络是基于我们提出的剩余局部和非局部注意块，其中包括主干和掩码分支。网络在掩码分支中通过非局部块获得非局部混合注意。这种注意机制有助于从层次特征中学习局部和非局部信息。
• 我们提出了保留更多低层特征的剩余非局部注意学习来训练深层网络，更适合于图像恢复。利用来自深层网络的非局部低层和高层注意，可以追求更好的网络表征能力，最终得到高质量的图像恢复结果。
• 我们通过大量实验证明，我们的RNAN对于各种图像恢复任务都是强大的。RNAN在图像去噪、去噪、压缩伪影减少和超分辨率方面取得了优于主流方法的效果。此外，RNAN在模型尺寸适中的情况下实现了优异的性能，并且执行速度非常快。

3 RESIDUAL NON-LOCAL ATTENTION NETWORK FOR IMAGE RESTORATION

3.1 FRAMEWORK

我们提出的剩余非局部注意网络（RNAN）的框架如图1所示。让我们将IL 和IH表示为低质量（例如，噪声、模糊或压缩图像）和高质量图像。重建图像IR可以通过下式获得：

其中HRNAN表示我们提出的RNAN的函数。通过在像素空间使用全局残差学习，我们网络的主要部分可以集中学习退化成分（例如，噪声、模糊或压缩的伪影）。

第一层和最后一层卷积分别是浅层特征提取层和重构层。我们提出了剩余的局部和非局部注意块来提取层次化的注意感知特征。除了让主网络学习退化分量外，我们还通过使用局部和非局部注意来进一步关注更具挑战性的领域。我们只在低层次和高层次的特征空间中加入剩余的非局部注意块。这主要是因为少数非局部模块可以很好地为网络提供非局部的图像恢复能力。

然后利用损失函数对RNAN进行优化。研究了几种损失函数，如L2，L1，感知和对抗性损失。为了证明我们的RNAN的有效性，我们选择了与以前工作相同的损失函数（如L2loss函数）进行优化。给定一个训练集，其中包含N个低质量输入及其高质量对应项。训练RNAN的目的是最小化L2loss函数：

如第4节所述，我们使用与其他比较方法相同的损失函数。这样的选择使得我们

所提出的RNAN的有效性更加清晰和公平。然后对剩余的局部和非局部注意块进行了详细的分析。

3.2 RESIDUAL NON-LOCAL ATTENTION BLOCK

我们的剩余非局部注意网络是由图2所示的几个残差局部和非局部注意块叠加而成的。每个注意块被分成两部分：在注意块的开始和结束处的q个剩余块（RBs）。中间有两个分支：主干分支和掩膜分支。对于非局部注意块，我们将非局部块（NLB）加入到mask分支中，得到非局部注意。然后我们给出这些组件的更多细节。

3.2.1 TRUNK BRANCH

如图2所示，主干分支包括t个剩余块（RBs）。与ResNet中的原始剩余区块（He et al.，2016）不同，我们采用了（Lim et al.，2017）中的简化RB。简化的RB（用蓝色虚线标记）仅由两个卷积层和一个ReLU（Nair&Hinton，2010）组成，省略了不必要的组件，例如maxpooling和batch normalization（Ioffe&Szegedy，2015）层。我们发现这种简化的RB不仅有助于图像的超分辨率（Lim et al.，2017），而且有助于为其他图像恢复任务构建非常深的网络。

不同深度主干分支的特征图作为层次特征。如果不考虑注意机制，该网络将成为一个简化的ResNet。通过mask分支，我们可以利用通道和空间注意自适应地重缩放层次特征。然后我们给出了关于局部和非局部注意的更多细节。

3.2.2 MASK BRANCH

如图2中红色虚线所示，我们网络中使用的掩码分支包括局部和非局部分支。在这里，我们主要关注局部掩模分支，它可以通过使用非局部块（NLB，用绿色虚线箭头标记）变成非局部掩模分支。

掩模分支的关键是如何掌握更大范围的信息，即更大的感受野大小，从而获得更复杂的注意图。一种可能的解决方案是多次执行maxpooling，如（Wang et al.，2017）中用于图像分类的方法。然而，在图像恢复中需要更精确的像素级结果。Maxpooling会丢失图像的很多细节，导致性能下降。为了克服这些缺点，我们选择大步长卷积和反卷积来扩大感受野的大小。另一种方法是在整个输入中考虑非局部信息，这将在下一小节中讨论。

从输入，stride≥ 2的卷积层增加了m个RBs后的感受野大小。在额外的2m个RBs之后，缩小后的特征图由反卷积层（也称为转置卷积层）展开。上采样后的特征通过m个RBs和1个1× 1个卷积层进一步前传。然后一个sigmoid层将输出值标准化，范围为[0,1]。虽然掩膜支的感受野比主干支大得多，但不能一次覆盖全部特征。这可以通过使用非局部块（NLB）来实现，从而产生非局部混合注意。

3.2.3 NON-LOCAL MIXED ATTENTION

如上所述，卷积运算一次处理一个局部邻域。为了获得更好的注意力图，这里我们试图一次考虑所有的位置。受经典的非局部平均方法和非局部神经网络的启发，我们将非局部块（NLB）并入掩码分支以获得非局部混合注意（如图3所示）。非局部操作可以定义为

其中i是输出特征位置索引，j是枚举所有可能位置的索引。x和y是非局部操作的输入和输出。成对函数计算Xi和Xj之间的关系。函数g（xj）计算位置j处输入的表示。

如图3所示，我们使用嵌入的高斯函数来评估成对关系

其中Wu和Wv是权重矩阵。如（Wang等人，2018a）中所研究的，f有几种版本，例如高斯函数、点积相似性和特征串联。对于g我们还考虑了加权矩阵为Wg的线性嵌入：。然后在非局部块（NLB）的位置i处的输出z被计算为

其中Wz是一个权重矩阵。对于给定的i，方程3中的变成了沿维度j的softmax。剩余连接允许我们将Wz初始化为0来将NLB插入预训练网络（Wang等人，2018a）。

通过非局部和局部注意计算，最终用sigmoid函数映射掩模分支中的特征映射

其中i范围在空间位置上，c范围在特征通道位置上。这种简单的sigmoid操作应用于每个通道和空间位置，导致混合注意力（Wang等人，2017）。因此，具有非局部块的掩模分支可以产生非局部混合注意。然而，简单地将主干和掩码分支的特征相乘并不能形成很深的可训练网络。我们提出了剩余非局部注意学习来解决这些问题。

3.3 RESIDUAL NON-LOCAL ATTENTION LEARNING

如何训练具有非局部混合注意的深度图像恢复网络还不清楚。这里我们只考虑主干和掩码分支，以及它们之间的残差连接（图2）。我们着重于从输入特征x中获取非局部注意信息。值得注意的是，Wang等人（2017）提出了一种形式的注意残差学习，其公式如下：

我们发现这种形式的注意学习不适合于图像恢复任务。这主要是因为等式7更适用于高级视觉任务（例如，图像分类），其中低级特征没有保留太多。然而，低层特征对于图像恢复更为重要。通过直接引入输入特征x，提出了一种简单而更适合的剩余注意学习方法。我们计算它的输出HRNA(x)为

其中Htrunk（x）和Hmask（x）分别表示主干和掩码分支的函数。这样的残差学习往往能保留更多的低层特征，使我们能够形成非常深的网络，用于具有更强表示能力的高质量图像恢复任务。

3.4 IMPLEMENTATION DETAILS

现在，我们详细说明了我们提出的RNAN的实现细节。我们使用10个剩余的局部和非局部注意块（2个非局部注意块）。在每个残差（非）局部块中，我们设置q，t，m=2，2，1。我们设置3×3作为所有卷积层的大小，除了非局部块和sigmoid函数之前的卷积层，它们的核大小为1×1.RBs中的特征有64个过滤器，除了非局部块（见图3），其中C=32。在每个训练批中，大小为48×48 的16个低质量（LQ）补丁作为输入。我们的模型是由ADAM optimizer用β1= 0.9, β2=0.999，=10−8.初始学习速率设置为10−4然后每200000次反向传播迭代减小到一半。我们在Titan Xp GPU上使用PyTorch实现了所有的模型。

4 EXPERIMENTS

我们将我们提出的RNAN应用于三个经典的图像恢复任务：图像去噪、去马赛克和压缩伪影。对于图像去噪和去马赛克，我们遵循与IRCNN相同的设置（Zhang等人，2017b）。对于图像压缩伪影的减少，我们遵循与ARCNN相同的设置（Dong等人，2015）。我们在DIV2K中使用了800张训练图像（Timofte et al.，2017；Agustsson&Timofte，2017）训练我们的所有模型。对于每项任务，我们使用常用的数据集进行测试，并报告PSNR和SSIM（Wang等人，2004）来评估每种方法的结果。更多结果见附录A。

4.1 ABLATION STUDY

我们在表1中展示了消融研究，以研究RNAN中不同组件的影响。

4.2 COLOR AND GRAY IMAGE DENOISING

我们将RNAN与最先进的去噪方法进行比较：BM3D（Dabov et al.，2007b）、CBM3D（Dabov et al.，2007a）、TNRD（Chen&Pock，2017）、RED（Mao et al.，2016）、DnCNN（Zhang et al.，2017a）、MemNet（Tai et al.，2017）、IRCNN（Zhang et al.，2017b）和FFDNet（Zhang et al.，2017c）。Kodak24(http://r0k.us/graphics/kodak/)BSD68（Martin等人，2001）和Urban100（Huang等人，2015）用于彩色和灰度图像去噪。不同级别的AWGN噪声（例如，10、30、50和70）被添加到干净的图像中。

定量结果如表2和表3所示。我们可以看到，我们提出的RNAN在所有噪声水平下的所有数据集上都达到了最佳结果。我们提出的非局部注意覆盖了整个图像的信息，这对于去除重图像噪声是有效的。为了证明这一分析，我们采用了噪声级σ = 70为例。我们可以看到，我们提出的RNAN比第二好的方法FFDNet获得了0.48、0.30和1.06db的PSNR增益。这种比较有力地证明了我们提出的非局部混合注意的有效性。

我们也在图4和图5中展示了视觉结果。通过学习非局部混合注意，RNAN区别对待不同的图像部分，明显减轻了过度平滑的伪影。

4.3 IMAGE DEMOSAICING

按照IRCNN（Zhang等人，2017b）中的相同设置，我们将图像去马赛克结果与IRCNN在McMaster（Zhang等人，2017b）、Kodak24、BSD68和Urban100上进行了比较。由于IRCNN是用于图像去马赛克和有限空间的最佳方法之一，我们在表四中仅与IRCNN进行比较。正如我们所看到的，马赛克图像的质量非常差，导致PSNR和SSIM值非常低。IRCNN可以对低质量图像进行增强，获得较高的PSNR和SSIM值。与IRCNN相比，我们的RNAN仍然可以做出显著的改进。利用局部和非局部注意，我们的RNAN可以更好地处理退化情况。

视觉效果如图6所示。尽管IRCNN可以大大消除马赛克效果，但其结果中仍然存在一些伪影（例如，“img 026”中的块效应）。然而，RNAN恢复更可靠的颜色并缓解块效应。

4.4 IMAGE COMPRESSION ARTIFACTS REDUCTION

我们进一步应用RNAN来减少图像压缩伪影。我们将RNAN与SA-DCT（Foi et al.，2007）、ARCNN（Dong et al.，2015）、TNRD（Chen&Pock，2017）和DnCNN（Zhang et al.，2017a）进行了比较。我们采用标准JPEG压缩方案，通过以下方式获得压缩图像（Dong等人，2015）。Matlab JPEG编码器中使用了四种JPEG质量设置q=10、20、30、40。在这里，我们只关注Y通道的恢复（在YCbCr空间中），以保持与其他方法的公平比较。我们在ARCNN中使用相同的数据集LIVE1（Sheikh et al.，2005）和Classic5（Foi et al.，2007），并在表5中报告PSNR/SSIM值。如我们所见，我们的RNAN在LIVE1和Classic5上实现了最好的PSNR和SSIM值。

我们在图7中进一步展示了视觉比较。我们提供了非常低的图像质量（q=10）下的比较。块效应可以在一定程度上消除，但是ARCNN、TNRD和DnCNN也会使某些结构过度平滑。RNAN通过考虑非局部混合注意，获得结构一致的更多细节。

4.5 IMAGE SUPER-RESOLUTION

我们进一步将我们的RNAN与最先进的SR方法进行比较：EDSR（Lim et al.，2017）、SRMDNF（Zhang et al.，2018a）、D-DBPN（Haris et al.，2018）和RCAN（Zhang et al.，2018b）。类似于（Lim等人，2017；Zhang等人，2018c），我们还引入了自集成策略来进一步改进我们的RNAN，并将自集成策略表示为RNAN+。

如表6所示，我们的RNAN+在基准数据集中取得了第二好的性能：Set5（Bevilacqua et al.，2012）、Set14（Zeyde et al.，2010）、B100（Martin et al.，2001）、Urban100（Huang et al.，2015）和Manga109（Matsui et al.，2017）。即使没有自集成，我们的RNAN在大多数情况下也能获得第三好的结果。这种改进是显著的，因为RNAN的参数数为7.5 M，远小于EDSR中的43 M和RCAN中的16 M。我们的RNAN（约120个卷积层）的网络深度也比RCAN（约400个卷积层）浅得多。结果表明，非局部注意力能更好地利用主网，节省大量网络参数。

在图8中，我们执行图像SR(×4） 用几种最先进的方法。我们可以看到我们的RNAN在结构更精细的情况下获得了更好的视觉效果。这些比较进一步证明了我们提出的RNAN在使用非局部混合注意时的有效性。

4.6 PARAMETERS AND RUNNING TIME ANALYSES

我们还比较了表7中基于彩色图像去噪的参数、运行时间和性能(σ=50). 具有10个块的RNAN以最大的参数个数获得最佳性能，可以减少到只有2个块，获得次优性能。在这里，我们报告运行时间以供参考，因为时间与实现平台和代码有关。

5 CONCLUSIONS

本文提出了一种用于高质量图像恢复的残差非局部注意网络。该网络通过叠加局部和非局部注意块来构建，提取局部和非局部注意感知特征，由主干和（非）局部掩码分支组成。它们用于提取层次特征，并使用软权值自适应地重新缩放层次特征。通过考虑整个特征图，进一步产生非局部注意。此外，我们还提出了剩余局部和非局部注意学习来训练深层网络。在注意计算中引入了输入特征，更适合于图像恢复。RNAN以适中的模型尺寸和运行时间实现了最先进的图像恢复结果。