ELAN:将超分网络SwinIR高效化,最快可达4.5倍

2022-09-01 11:28:21 浏览数 (3)

关注公众号,发现CV技术之美

本篇分享论文『Efficient Long-Range Attention Network for Image Super-resolution』,由香港理工大学、OPPO 提出ELAN:将超分网络SwinIR高效化,最快可达4.5倍。

  • 作者单位:香港理工大学、OPPO
  • 论文链接:https://arxiv.org/abs/2203.06697
  • 项目链接:https://github.com/xindongzhang/ELAN

01 看点

Transformer通过利用自注意(self-attention, SA)进行特征提取,取得了令人印象深刻的结果。然而SA的计算是非常昂贵的,并且一些操作对于SR任务来说可能是冗余的。本文提出了一种高效的远程注意网络(ELAN),它首先采用shift卷积在保持与1x1卷积相当的复杂度情况下有效的提取图像局部结构信息,然后提出了一种分组多尺度自注意(GMSA)模块,它使用不同的窗口大小在非重叠的特征组上计算SA。通过将两个shift卷积与GMSA模块级联,构建高效的远程注意块(ELAB),并通过共享注意机制进一步加速该模块的运行,极限情况下能比SwinIR快四倍!

02 方法

Overview

ELAN的架构如下图所示,它由三部分组成:浅层特征提取、深度特征提取和HR图像重建。其中浅层特征提取由一个3x3卷积组成,深度特征提取由堆叠的ELAB和一个残差连接组成,重建模块由一个3 × 3卷积和一个PixelShuffle操作组成,损失采用范式。

ELAB

如下图所示,ELAB由局部特征提取和GMSA组成,上述都配备了残差连接。

局部特征提取:中间特征以往多采用多层感知或两个1 × 1卷积的方法,但这只有1 × 1的感受野。本文采用中间有ReLU激活函数的两个shift卷积来扩大感受野以更有效地提取局部特征。shift卷积由shift运算和1 × 1卷积组成。具体的说,shift卷积将输入特征平均分为五组,前四组特征沿不同的空间维度进行位移,最后一组不变,然后用1x1卷积就可以利用位移后的相邻元素的信息。这没有引入额外的可学习参数并保持了与1x1卷积相似的复杂度,同时使感受野由1变成了3。

GMSA: 给定一个CxHxW的特征图,使用窗口大小为M的自注意的计算复杂度为

。GMSA首先将输入特征分成K组,然后利用

的窗口大小计算第K组特征的SA。假设通道平均分割且,单个组的计算复杂度为

,总复杂度为

,然后将不同组的SA输出通过1x1卷积聚合。

加速SA(ASA):首先,摒弃了以往Transformer中广泛使用的层归一化(LN),因为LN将SA的计算分割成许多元素级操作,这不利于高效推理。本文采用批归一化(BN)来稳定训练过程,在推理阶段BN可以合并到卷积运算中,不会造成额外的计算成本。其次,SwinIR中的SA使用三个独立的1 × 1卷积θ、φ和g计算,本文设置θ = φ,计算对称高斯空间中的SA,这可以节省一个1 × 1卷积但不牺牲SR的性能,如上图(e)所示。

共享注意得分图: 上述SA的一次前向传递包含两个1 × 1卷积和四个reshape操作。由于SR任务中的特征尺寸较大,reshape耗时较长。为此,本文在相邻的SA模块之间共享注意得分图。如下图b所示,第i个SA模块的注意得分图,直接被后面n个SA模块重复使用。这种方法在n个SA时减少2n次reshape和n个1×1卷积。实验发现,使用少量的n(如1或2)时只会导致SR性能的轻微下降,但节省了大量的计算资源。

移位窗口:本文改进了SwinIR的移位窗口机制,如上图c所示。首先对特征进行对角线方向的循环偏移,并计算偏移后的GMSA。然后将结果反向循环偏移。利用半窗口大小的循环偏移可以对特征映射进行新的划分,并在之前的GMSA模块中引入相邻非重叠窗口之间的连接。虽然循环偏移将边界上的一些像素点转移到较远的区域,但这些像素点在SR任务中只占小部分,所以循环偏移对SR的影响很小。利用循环移位机制就可以去掉SwinIR中所采用的掩蔽策略和相对位置编码,使得网络更加整洁高效。

01 实验

消融实验

移位窗口机制,推断速度从247ms降低到177ms的同时性能几乎保持不变。用ASA来替代SA,在不损失PSNR/SSIM性能的情况下,推理延迟从177ms降低到66ms。通过采用GMSA对长期依赖模型进行有效建模,PSNR和SSIM指数在所有5个数据集上都得到了显著改善。这表明GMSA比SwinIR小窗口下的SA更有效。共享注意机制,可以在性能下降很小的情况下,进一步加快光的推断时间。

定量评估

轻量级SR模型对比Transformer的方法通过利用图像的自相似性,在PSNR/SSIM指标上优于许多基于cnn的方法。但是因为在SwinIR中SA是一个沉重的计算负担,SwinIR-light的延迟比CNN的方法要慢×10以上。得益于高效的远程注意力设计,ELAN-light模型不仅在所有五个数据集上获得了较好的指标,而且比SwinIR-light快了×4.5。同时参数和计算复杂度也比SwinIR-light要小。

经典SR模型对比ELAN在所有数据集上都获得了更好的PSNR和SSIM,推理速度×2。

定性评估

0 人点赞