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

关注公众号，发现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。

定性评估