【Transformer+SR】ESRT：图像超分中的超轻量Transformer

arXiv:https://arxiv.org/pdf/2108.11084.pdf

Abstract

SISR见证了深度学习的迅速发展，然而现有大多方法聚焦于构建更复杂、更深的网络，进而带来大量的计算复杂度与内存占用。近来，NLP领域的Transformer开始在CV任务上发力并逐步占领了不少领域的“高地”。但是，Transformer的高计算复杂度与高GPU占用导致其难以设计的过深。

为解决上述问题，本文提出了提出快速而精确的图像超分方案ESRT(Efficient Supe-Resolution Transformer)，它是一种CNN-Transformer混合架构。ESRT包含两个骨干模块：LCB(lightweight CNN Backbone)与LTB(lightweight Transformer Backbone)，其中LCB以更低的计算量提取深层特征并动态调整特征的尺寸；LTB则由多个高效Transformer模块(EMHA)构成。在EMHA内部，FSM(Feature Split Module)用于将长序列拆分为多个子序列，然后在这些子序列上执行注意力操作。该模块可以大幅降低GPU内存占用。

实验结果表明：所提ESRT可以取得极具竞争力的结果。相比原始Transformer(GPU占用为16057M)，所提ET的GPU占用仅为4191M，同时具有更好的性能。

Method

接下来，我们先对所提ESRT的整体架构进行介绍；然后，我们提出带HPB(high preserving block)与HFM(high frequency filtering module)的LCB；其次，我们提出了高效Transformer模块LTB；最后我们对ESRT与其他方案的差异进行了对比。

Overall Architecture

上图给出了所提ESRT整体架构图，它包含四个主要部件：

Lightweight CNN Backbone(LCB)

类似于其他超分模型，LCB由多个HPB模块构成并用于提取深层特征，此时模型已经具备初步的超分能力。

Adaptive Residual Feature Block 已有研究表明：网络的深度与性能高度相关，而残差学习有助于缓解梯度消失问题。因此，我们提出一种自适应残差特征模块，见下图。从中可以看到：ARFB包含两个RU与两个卷积 。此外，它引入了类似AWSRN中的自适应残差学习机制。

High Preserving Block 下图给出了所提HPB模块结构示意图，它用于保持高频信息，HBP主要采用ARFB构建。为尽可能降低计算量，HPB中还进行了特征的下采样，将大部分计算量在更低分辨率特征上执行。

Lightweight Transformer Backbone(LTB)

在这里，我们旨在设计一种轻量型的Transformer模块以捕获长距离依赖关系(类似非局部相似性)。

Pre- and Post-process for ET 一般来讲，标准Transformer采用1D序列作为输入，学习序列之间的长距离依赖关系。然而，CV任务的输入往往为2D图像，一种常见操作：通过像素重拍将2D图像reshape为1D序列。但这种序列化操作会破坏图像的局部相关性，导致次优性能。

The Architecture of ET 下图给出了本文所提ET的架构示意图。为简单性与高效性，ET仅采用了编码器结构：EMHA+MLP。

在EMHA中，我们进行了多个改进以使其更高效同时占据更少的GPU内存。假设输入嵌入的尺寸为。首先，采用Reduction层降低通道数，然后采用Linear层将特征投影为三个元素Q、K以及V（这里与标准Transformer类似）。

在SR任务中，由于输入分辨率可能非常大，采用原始的Transformer操作可能会占用过多的显存；此外，超分图像的预测结果往往仅依赖于局部近邻区域。因此，我们采用FSM将Q、K、V拆分为s份，这可以极大的降低计算量以及GPU显存占用。然后，再采用自注意力计算公式对每份的QKV进行注意力计算。

Implements Details

在所提ESRT中，除Reduction模块外，卷积的尺寸默认为。除了Fusion层外卷积层的通道数为32。在图像重建部分，我们采用PixelShuffle操作进行特征上采样。HFP中的k为2；HPB的模块数为3，LTB中的ET数量为1，ET中的拆分因子s为4，ET的预处理与后处理中的k为3，EMHA中的head数为8.

Experiments

上表给出了所提方案与其他轻量型方案的性能对比，从中可以看到：ESRT取得了最佳性能 。

来源：AIWalker
作者： Happy