手把手教你改进PAN！董超团队开源超大感受野注意力超分方案VapSR

会员服务 ·

手把手教你改进PAN！董超团队开源超大感受野注意力超分方案VapSR

2022 年 10 月 24 日 极市平台

↑ 点击蓝字关注极市平台

作者丨HappyAIWalker

来源丨AIWalker

编辑丨极市平台

极市导读

注意机制在设计高级超分辨率（SR）网络中起着关键作用。在这项工作中，作者们通过改进注意机制设计了一个高效的SR网络。VapSR以更少的参数优于当前的轻量级网络。>>加入极市CV技术交流群，走在计算机视觉的最前沿

paper：https://arxiv.org/abs/2210.05960

code：https://github.com/zhoumumu/VapSR

注意力机制是深度学习领域非常重要的一个研究方向，在图像超分领域也有不少典型的应用案例，比如基于通道注意力构建的RCAN，基于二阶注意力机制构建的SAN，基于像素注意力机制构建的PAN，基于Transformer自注意力机制构建的SwinIR，基于多尺度大核注意力的MAN等。

本文则以PAN为蓝本，对其进行逐步改进以期达到更少的参数量、更高的超分性能。该方案具体包含以下几个关键点：

提升注意力分割的感受野，类似大核卷积注意力VAN；
将稠密卷积核替换为深度分离卷积，进一步降低参数量；
引入像素规范化(Pixel Normalization)技术，其实就是Layer Normalization，但出发点不同。

上述关键技术点为注意力机制的设计提供了一个清晰的演变路线，最终得到了本文的VapSR，即大感受像素注意力网络(VAst-receptive-field Pixel attention Network)。

实验结果表明：相比其他轻量超分网络，VapSR具有更少的参数量。比如，项目IMDB与RFDN，VapSR仅需21.68%、28.18%的参数即可取得与之相当的性能。

本文动机

通过引入像素注意力，PAN在大幅降低参数量的同时取得了非常优秀的性能。相比通道注意力与空域注意力，像素注意力是一种更广义的注意力形式，为进一步的探索提供了一个非常好的基线。

受启发于自注意力的发展，我们认为：基于卷积操作的注意力仍有进一步改进的空间。因此，作者通过以下三个像素注意力中的设计原则展示了改善超分注意力的过程：

首先，在注意力分支引入大核卷积具有明显的优势；
其次，深度分离卷积可以降低大核卷积导致的巨大计算复杂度问题；
最后，引入像素规范化操作让训练更高效、更稳定。

Large Kernel 以上图i中的baseline为基础，作者首先对注意力分支进行感受野扩增：将提升到(将图示ii)，性能提升0.15dB，但参数量从846K提升到了4123K。

Parameter Reduction 为降低参数量，我们尝试尽可能移除相对不重要的部分。作者提出了三个方案：(1) 将非注意力分支的卷积尺寸从下调到；(2) 将大核卷积注意力分支替换为深度深度分离卷积；(3) 将深度分离卷积中的深度卷积进行分解为深度卷积+带扩张因子的深度卷积(该机制可参考下图，将卷积拆分为+，其中后者的扩张因子为3)。此时，模型性能变为28.48dB，但参数量降到了240K，参数量基本被压缩到了极限。

Pixel Normalization(PN) 注意力机制的元素乘操作会导致训练不稳定问题：小学习率收敛不够好，大学习率又会出现梯度异常。前面的注意力改进导致所得方案存在性能下降问题。为解决该问题，作者经深入分析后提出了像素规范化技术(可参考下图不同规范化技术的可视化对比)。

假设输入特征为 , 第个像素的特征均值与方差可以描述如下:

那么, 像素规范化可以表示为：

那么，像素规范化可以表示为：当引入PN后，模型的性能取得了显著的提升，达到了28.92dB，参数量仅为241K。

Switch Attention to Middle 在上述基础上，作者进一步将注意力的位置进行了调整，放到了两个卷积中间。此时，模型性能得到了0.03dB提升，达到了28.95dB，参数量仍为241K。

本文方案

前面的探索主要聚焦在微观层面，基于此，作者进一步在宏观层面进行了更多设计与提炼，进而构建了VapSR，取得了更佳的性能，同时具有更少的参数量。

上图给出了所提VapSR架构示意图，延续了常规轻量方案的设计思路：

浅层特征: ;
非线性映射:

图像重建:

模块在前面探索得到的模块上进行了微调: (1) 主要是将模块输入与输出通道数从 64 减少到了 48 , 保持中间注意力部分的通道数仍为 (2) 将注意力分支深度扩张卷积 (有时也称之为空洞卷积) 调整为深度扩张卷积, 此时感受野为 (3) 调整了注意力分支三个卷积的顺序, 将卷积移到最前面。对于 VapSR-S, 作者进一步将部分从卷积调整为组卷积 , 该操作可以进一步降低参数量。

class Attention(nn.Module):  
    def __init__(self, dim):  
        super().__init__()  
        self.pointwise = nn.Conv2d(dim, dim, 1)  
        self.depthwise = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)  
        self.depthwise_dilated = nn.Conv2d(dim, dim, 5, 1, padding=6, groups=dim, dilation=3)  

    def forward(self, x):  
        u = x.clone()  
        attn = self.pointwise(x)  
        attn = self.depthwise(attn)  
        attn = self.depthwise_dilated(attn)  
        return u * attn  

class VAB(nn.Module):  
    def __init__(self, d_model, d_atten):  
        super().__init__()  
        self.proj_1 = nn.Conv2d(d_model, d_atten, 1)  
        self.activation = nn.GELU()  
        self.atten_branch = Attention(d_atten)  
        self.proj_2 = nn.Conv2d(d_atten, d_model, 1)  
        self.pixel_norm = nn.LayerNorm(d_model)  
        default_init_weights([self.pixel_norm], 0.1)  

    def forward(self, x):  
        shorcut = x.clone()  
        x = self.proj_1(x)  
        x = self.activation(x)  
        x = self.atten_branch(x)  
        x = self.proj_2(x)  
        x = x + shorcut  
      
        x = x.permute(0, 2, 3, 1) #(B, H, W, C)  
        x = self.pixel_norm(x)  
        x = x.permute(0, 3, 1, 2).contiguous() #(B, C, H, W)  
      
        return x