【Transformer+SR】ESRT：图像超分中的超轻量Transformer
【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整体架构图，它包含四个主要部件：
- 浅层特征提取：它采用卷积从输入图像中提取浅层特征
- 轻量CNN骨干LCB：它对前述所提取的特征进行深层特征提取，可以描述如下：
- 轻量Transformer骨干LTB：将上述CNN提取的所有特征拼接后送入到LTB进行融合，描述如下：
- 图像重建模块：为得到超分结果，我们将前述特征送入到重建模块，描述如下：
Lightweight CNN Backbone(LCB)
类似于其他超分模型，LCB由多个HPB模块构成并用于提取深层特征，此时模型已经具备初步的超分能力。
High-frequency Filtering Module 在介绍HPB之前，我们先介绍HFM，见上图。由于傅里叶变换难以嵌入到CNN中，故我们提出一种微分HFM，其目标在于估计图像的高频信息。假设输入特征
的尺寸为
，我们首先采用AVP操作
，然后上采样得到新的张量
,最后我们采用减操作提取高频信息。下图给出了特征激活示意图，可以看到：
更为平滑.
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序列。但这种序列化操作会破坏图像的局部相关性，导致次优性能。
为处理2D特征，我们采用unfolding操作将特征拆分为块，每个块视作一个word。具体来说，特征
将被unfolded为
。注：由于unfolder操作能够反应每个块的位置信息，这里并没有引入位置嵌入信息。这是真的吗？保持怀疑态度 。
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取得了最佳性能 。
吐槽时间
看到一半的时候就真的看不下去了，ESRT的架构人工设计痕迹太明显：各种组合、各种复杂 。尤其要吐槽的是LCB这里，轻量是轻量了，但推理速度看着就慢啊！LCB这里能看到各种思想的组合，比如源自AWSRN的自适应学习、ARFB好像也是某一些文章里面的设计模块，
这篇文章唯一有点价值的反而是ET，不过这里的unfolded不是很常规的操作吗？各种最近2-3个月出的Transformer不都是这样去设计的吗？当然，ET中结合超分任务特性的拆分降低计算量是值得表扬的。