前言
论文:https://arxiv.org/pdf/2203.06697.pdf
代码:https://github.com/xindongzhang/ELAN
一、Introduction
最近,基于Transformer的方法已经证明了令人印象深刻的效果,在各种视觉任务中,包括图像超分辨率(SR),利用自我注意(SA)的特征提取。然而,在大多数现有的基于Transformer的模型中,SA的计算是非常昂贵的,而一些所采用的操作对于SR任务来说可能是冗余的。这限制了SA计算的范围并且因此限制了SR的性能。
在这项工作中,本文提出了一种用于图像超分的高效长距离网络ELAN(Efficient Long-Range Attention Network)。具体来说,我们首先采用移位卷积(Shift Convolution)提取图像的局部结构信息同时保持与 1x1 卷积相当的复杂度;然后提出了一种GMSA(Group-wise Multi-scale Sel-Attention)模块,它在不同窗口尺寸特征上计算自注意力以探索长程依赖关系。通过简单地将两个移位卷积与一个GMSA模块级联,构建了一个高效的远程注意块(ELAB),并利用共享注意机制进一步加速了该模块的运行。实验结果表明:相比于其他Transformer方案,所提的ELAN可以取得更佳结果,同时具有更低的复杂度。
二、Method
1、Network Architecture
ELAN的整体流程图如图1(a)所示,该流程由三个模块组成:浅层特征提取、基于ELAB的深度特征提取和HR图像重建。该网络的拓扑结构比较简单,从浅层特征提取模块输出到深度特征提取模块输出之间有一个全局的快捷连接,然后再馈送给HR重构模块。
(a)ELAN的整体流程图,包含几个ELABs,两个3x3卷积核一个pixel shuffle算子。
如图1(b)所示,我们的ELAB由局部特征提取模块(Local Feature Extraction)和组多尺度注意力(Group-wise Multi-scale Sel-Attention,GMSA)模块组成,两者都配备了残差学习策略。
(b) 高效远程注意力块(ELAB)的架构。 (c) 移位卷积的图示,它由移位运算和一个 1 × 1 卷积组成。 (d) 分组多尺度自注意力(GMSA)的计算说明。 (e) 我们的加速自注意力(ASA)计算的图示:
2、Local feature extraction
给定中间特征X,之前的研究大多通过多层感知或两个级联的1×1卷积来提取局部特征,然而它们只有1×1的感受野。为了扩大感受野以更有效地提取局部特征,我们利用两个移位卷积,并在它们之间进行简单的 ReLU 激活。如图 1(c) 所示,shift-conv 由一组移位操作和一个 1×1 卷积组成。
具体来说,我们将输入特征平均分为五组,并将前四组特征沿不同的空间维度移动,包括左、右、上、下,而最后一组保持不变。因此,接下来的 1 × 1 卷积可以利用相邻像素的信息。在不引入额外的可学习参数和大量计算的情况下,shift-conv 可以提供更大的感受野,同时保持与 1 × 1 卷积几乎相同的算术复杂度。
3、Group-wise multi-scale self-attention (GMSA)
为了更有效地计算远程 SA,我们提出了 GMSA 模块,如图 1(d)所示。我们首先将输入特征 X 分成 K 组,用 表示,然后使用窗口大小 Mk 计算第 k 组特征上的 SA。这样,我们就可以通过设置不同窗口大小的比例来灵活控制计算成本。例如,假设将K组特征均分到C K 个通道,则K组SA的计算成本为 。然后,在不同组上计算的 SA 通过 1 × 1 卷积连接并合并。
4、Accelerated self-attention (ASA)
ASA 如图 1(e) 所示。在现有的Transformer模型中,SA 的计算是计算和内存密集型的。我们做了一些修改来加速 SA 的计算,特别是在推理阶段。首先,我们丢弃了之前的 Transformer 模型中广泛采用的层归一化(LN),因为 LN 将 SA 的计算分割为许多元素级操作,这对于高效来说并不友好推理。相反,我们利用批量归一化(BN)来稳定训练过程。值得一提的是,BN可以合并到卷积运算中,这不会在推理阶段造成额外的计算成本。。其次,SwinIR [28]中的SA是在嵌入的高斯空间上计算的,其中使用三个独立的1×1卷积(用θ、ψ和g表示)将输入特征X映射到三个不同的特征图。
5、Shared Attention
尽管重塑操作是无 FLOPs 的,但由于 SR 任务中的特征尺寸较大,该操作非常耗时。为了进一步加速整个网络的 SA 计算,我们建议在相邻 SA 模块之间共享注意力分数。如图2(b)所示,第i个SA模块中计算出的注意力分数直接被后续n个相同规模的SA模块重复使用。这样,我们就可以避免接下来的 n 个 SA 的 2n 次重塑和 n 1×1 卷积操作。我们发现,所提出的共享注意力机制仅通过使用少量的 n(例如 1 或 2)而导致 SR 性能略有下降,同时在推理过程中节省了大量的计算资源。
6、Shifted Window
在分组多尺度窗口上计算的 SA 仍然缺乏相同尺度内本地窗口之间的连接。我们改进了 SwinIR的移位窗口机制,为 SR 任务提供了一个简单而有效的移位方案。整个过程如图 2(c) 所示。
我们首先沿对角线方向对特征进行循环平移,并计算平移后的特征的 GMSA。然后我们通过逆循环移位将结果移回原处。半窗口大小的循环移位导致了特征图的新划分,并在先前的 GMSA 模块中引入了相邻非重叠窗口之间的连接。尽管边界上的一些像素通过循环移位移动到较远的区域,但我们发现它对最终 SR 性能的影响可以忽略不计,因为这些像素在 SR 任务中只占据整个特征图的一小部分。
受益于循环移位机制,我们删除了 SwinIR中采用的掩蔽策略和相对位置编码,使我们的网络更整洁、更高效。
2. Experiments
Comparison with classic SR models
为了验证ELAN的可扩展性,我们进一步将ELAN的普通版本与最先进的经典的面向性能的SR模型进行比较,包括EDSR,SRFBN,RNAN,RDN、RCAN、SAN、IGNN HAN、NLSA 和 SwinIR 。请注意,这些模型要么采用具有大量通道的非常深的网络拓扑,要么采用复杂的自注意力和非局部策略。
Quantitative comparison
Qualitative comparison
与多种超分方案进行了对比,从中可以看到:
- 在所有尺度与基准数据集上,ELAN均取得了最佳性能;
- 相比SwinIR,ELAN取得了最佳PSNR与SSIM,同时具有更少参数量与计算量、更快推理速度;
- 相比CNN方案,ELAN在重建性能方面局具有显著优势,这得益于大感受野与长程特征相关性建模;
- 尽管HAN与NLSA通过探索注意力机制与更深网络课要取得相当性能,但其计算量与内存负载非常昂贵。比如NLSA甚至无法在NVIDIA 2080Ti GPU上执行x2超分任务,而ELAN则可以凭借更少计算量取得比这些复杂方案更高的性能。
3. Conclusion
在本文中,我们提出了一种用于单图像超分辨率的高效远程注意网络(ELAN)。
ELAN 具有简洁的拓扑结构,具有顺序级联的高效远程注意力块(ELAB)。每个ELAB由具有两个顺序移位卷积的局部特征提取模块和分组多尺度自注意(GMSA)模块组成,以逐渐增加自注意(SA)的感受野。受益于我们的加速 SA 计算和共享注意力机制,ELAB 可以以非常有效的方式有效捕获局部结构和远程依赖关系。
大量的实验表明,在轻量级和性能导向的设置上,ELAN 可以比以前最先进的 SR 模型获得极具竞争力的性能,同时比以前基于Transformer的 SR 方法更加经济。
尽管我们的 ELAN 比 SwinIR 实现了显着的加速,但与那些基于 CNN 的轻量级模型相比,SA 的计算仍然是计算和内存密集型的。未来,我们将进一步探索 SA 的更有效实现或近似,以实现更多低级视觉任务。