Residual non-local attention networks for image restoration

mytzs123

已于 2024-04-23 09:12:10 修改

阅读量52

点赞数

分类专栏：图像增强文章标签：计算机视觉深度学习

于 2022-07-25 12:48:49 首次发布

本文链接：https://blog.csdn.net/mytzs123/article/details/125972568

版权

图像增强专栏收录该内容

31 篇文章 1 订阅

订阅专栏

本文提出了一种名为剩余非局部注意网络（RNAN）的深度学习模型，用于高质量图像恢复任务，如去噪、去马赛克和压缩伪影减少。RNAN解决了传统方法中局部卷积运算忽视图像信息不均匀分布的问题，通过设计局部和非局部注意块，捕获像素间的长期依赖，并自适应地调整特征。非局部注意块用于获取全局上下文，而残差学习则增强了网络的表示能力。实验表明，与现有方法相比，RNAN在定量和视觉效果上表现出色。

摘要由CSDN通过智能技术生成

用于图像恢复的剩余非局部注意网络

ABSTRACT

在本文中，我们提出了一种用于高质量图像恢复的剩余非局部注意网络。由于没有考虑图像中信息的不均匀分布，以前的方法受到局部卷积运算以及空间和通道特征的同等处理的限制。为了解决这个问题，我们设计了局部和非局部注意块来提取特征，捕捉像素之间的长期依赖关系，并更加关注具有挑战性的部分。具体来说，我们在每个（非）局部注意块中设计主干分支和（非）局部掩码分支。主干分支用于提取层次特征。局部和非局部掩模分支旨在混合注意力自适应地重新缩放这些层次特征。局部掩码分支专注于具有卷积运算的更多局部结构，而非局部注意力更多地考虑整个特征映射中的长程依赖。此外，我们提出了剩余局部和非局部注意力学习来训练深度网络，从而进一步增强了网络的表示能力。我们提出的方法可以推广到各种图像恢复应用，如图像去噪、去马赛克、压缩伪影减少和超分辨率。实验表明，与最近领先的方法相比，我们的方法在定量和可视化方面取得了相当或更好的结果

现有方法存在的问题：

然而，现有的基于CNN的上述方法主要存在三个问题。首先，这些网络的感受野大小相对较小。它们大多采用卷积运算局部提取特征，无法捕捉整个图像中像素之间的长期依赖关系。更大的感受野可以更好地利用训练输入和更多的上下文信息。这将有助于捕捉LQ图像的潜在退化模型，尤其是当图像遭受严重损坏时。其次，这些网络的独特能力也有限。让我们以图像去噪为例。对于带噪图像，噪声可能出现在平面和纹理区域。平面区域的噪声去除比纹理区域更容易。希望使去噪模型更加关注文本区域。然而，以往的去噪方法大多忽略了对噪声输入中不同内容的考虑，并对其进行了同等对待。这将导致输出过度平滑，一些纹理细节也无法恢复。第三，在这些网络中，所有信道特征都被同等对待。这种幼稚的处理方法在处理不同类型的信息（例如，低频和高频信息）方面缺乏灵活性。对于一组功能，一些功能包含更多与HQ image相关的信息，而其他功能可能包含更多与损坏相关的信息。为了更准确地恢复图像，应考虑通道之间的相互依赖性。

为了解决上述问题，我们提出了用于高质量图像恢复的深度剩余非局部注意网络（RNAN）。我们设计了剩余的局部和非局部注意块作为深度网络的基本构建模块。每个注意块由主干和掩码分支组成。我们为主干分支引入残差块（He等人，2016；Lim等人，2017），并提取层次特征。对于掩模分支，我们使用最大卷积和反卷积进行特征降尺度和升尺度，以扩大感受野大小。此外，我们在掩码分支中加入非局部块以获得剩余的非局部混合注意力。我们将RNAN应用于各种恢复任务，包括图像去噪、去马赛克和压缩伪影减少。大量实验表明，与其他最近领先的方法相比，我们提出的RNAN在所有任务中都达到了最先进的结果。据我们所知，这是第一次考虑图像恢复问题的剩余非局部注意力。

主要贡献：

我们提出了用于高质量图像恢复的深度残差非局部网络。强大的网络基于我们提出的剩余局部和非局部注意块，由主干和掩码分支组成。该网络通过掩码分支中的非局部块获得非局部混合注意。这种注意力机制有助于从层次特征中学习局部和非局部信息。

我们提出了剩余非局部注意力学习，通过保留更多的低层特征来训练深度网络，更适合于图像恢复。利用来自深度网络的非局部低层和高层注意力，我们可以追求更好的网络表示能力，最终获得高质量的图像恢复结果。

我们通过大量实验证明，我们的RNAN对于各种图像恢复任务都是强大的。RNAN在图像去噪、去马赛克、压缩伪影减少和超分辨率方面取得了优于主流方法的结果。此外，RNAN在模型尺寸适中的情况下实现了卓越的性能，并且执行速度非常快。

网络结构：

我们的剩余非局部注意网络是通过堆叠多个剩余局部和非局部注意块构建的，如图2所示。每个注意块分为两部分：在注意块的开始和结束处的q个剩余块（RBs）。中间有两个分支：主干分支和掩码分支。对于非局部注意块，我们在掩码分支中加入非局部块（NLB），从而产生非局部注意。然后，我们给出了这些组件的更多细节。

主干分支：

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

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

掩码分支：

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

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

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

代码：https://github.com/yulunzhang/RNAN/blob/master/DN_Gray

主要代码：

class NonLocalBlock2D(nn.Module):
    def __init__(self, in_channels, inter_channels):
        super(NonLocalBlock2D, self).__init__()
        
        self.in_channels = in_channels
        self.inter_channels = inter_channels
        
        self.g = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0)
        
        self.W = nn.Conv2d(in_channels=self.inter_channels, out_channels=self.in_channels, kernel_size=1, stride=1, padding=0)
        nn.init.constant(self.W.weight, 0)
        nn.init.constant(self.W.bias, 0)
        
        self.theta = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0)
        
        self.phi = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0)
 
    def forward(self, x):
 
        batch_size = x.size(0)
        
        g_x = self.g(x).view(batch_size, self.inter_channels, -1)
        
        g_x = g_x.permute(0,2,1)
        
        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
        
        theta_x = theta_x.permute(0,2,1)
        
        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
        
        f = torch.matmul(theta_x, phi_x)
       
        f_div_C = F.softmax(f, dim=1)
        
        
        y = torch.matmul(f_div_C, g_x)
        
        y = y.permute(0,2,1).contiguous()
         
        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
        W_y = self.W(y)
        z = W_y + x
 
        return z