Happy · 2022年10月24日 · 北京市

VapSR | 手把手教你改进PAN,董超团队提出超大感受野注意力超分方案,已开源!

image.png

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在大幅降低参数量的同时取得了非常优秀的性能。相比通道注意力与空域注意力,像素注意力是一种更广义的注意力形式,为进一步的探索提供了一个非常好的基线。

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

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

image.png

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

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

image.png

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

当引入PN后,模型的性能取得了显著的提升,达到了28.92dB,参数量仅为241K。

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

本文方案

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

image.png

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

本文实验

在实验部分,作者构建了VapSR与VapSR-S两个版本的轻量型超分方案:

  • VapSR:包含21个VAB模块,主干通道数为48;
  • VapSR-S:包含11个VAB模块,主干通道数为32。

此外,需要注意的是:对于X4模型,重建模块并未采用常规的轻量方案(Conv+PS),而是采用了类EDSR的重方案(Conv+PS+Conv+PS)。

image.png

image.png

image.png

上表&图给出了不同方案的性能与可视化效果对比,从中可以看到:

  • 所提VapSR取得了SOTA性能,同时具有非常少的参数量。
  • 在X4任务上,相比RFDN与IMDN,VapSR仅需21.68%/28.18%的参数量,即可取得平均0.187dB指标提升;
  • VapSR-S取得了与BSRN-S相当的性能,后者是NTIRE2022-ESR模型复杂度赛道冠军。
  • 在线条重建方面,VapSR具有比其他方案更精确的重建效果。
来源:AIWalker
作者:HappyAIWalker

推荐阅读

本文章著作权归作者所有,任何形式的转载都请注明出处。更多动态滤波,图像质量,超分辨相关请关注我的专栏AIWalker。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
6197
内容数
191
夯实深度学习知识基础, 涵盖动态滤波,超分辨,轻量级框架等
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息