FMEN | NTIRE2021-ESR最低内存占用超分方案

2022-06-24 14:12:30 浏览数 (1)

NTIRE2022各大竞赛结果已定,您上榜了吗?

今天介绍一篇NTIRE2022 Efficient Super Resolution竞赛方案,该方案取得了最低内存占用、第二快的推理速度。推理耗时与内存占用是EISR(Efficient Image Super Resolution)的两个重要考量因素。在该方向上比较知名的方案当属IMDN与RFDN,两者均采用split与concat对蒸馏与聚合的策略进行了探索。相反,序贯式(Sequential)结构避免了频繁的访问前置状态与额外节点,对降低内存占用与推理耗时更有益

延续该设计理念,本文设计了一种轻量型方案FMEN,它主要由高度优化的卷积与激活层堆叠而成,尽可能降低的特征融合模块的使用。本文提出了一种新颖的序贯式注意力分支,每个pixel根据局部与全局上下文被赋予不同的重要性因子以增强高频细节学习能力。此外,作者在EISR常用的残差模块基础上提出了ERB(Enhanced Residual Block)以加速模型推理速度

组合上述设计理念记得到了本文所提FMEN与FMEN-S。相比AIM2020-EISR冠军方案E-RFDN,所提FMEN-S推理速度快33%,内存占用降低74%(可参考下图);在NTIRE2022-EISR竞赛中,FMEN-S取得了内存占用最低、推理速度第二的成绩。

1Method

Memory Analysis

内存占用是EISR方案部署时非常重要的因素之一。对一个节点而言,其内存占用M包含四部分:输入特征M_{input} 、输出特征M_{output} 、kept特征M_{kept} (之前节点计算并在未来节点使用的特征)以及网络参数M_{net} (相对而言,该部分可忽略)。总体的内存占用可以表示如下:

M = M_{input} M_{output} M_{kept} M_{net}

本文以上述两种典型结构为参考进行内存分析(假设输入所占内存为1x):

  • 对于上图a所示结构,其峰值内存占用为:
M_{input} M_{output} approx 2times C times H times W
  • 对于上图b所示结构,其峰值存在占用为:
2 times N times C times H times W

(N表示参与特征融合的特征数量,上图中N=3).

也就是说,相比Plain结构,Fusion结构会占用更高的峰值内存

Network Architecture

直接将序贯式结构用于EISR并不是那么容易的:相比先进的Fusion结构,无论是直接式方案还是重参数机制辅助方案均存在性能下降问题。除了内存占用外,推理耗时是EISR方案的另一个重要因素。

为解决上述问题,基于序贯式结构,本文提出ERB与HFAB分别用于深层特征学习与特征增强,不仅可降低内存占用,同时可加速推理效率。本文最终所得结构FMEN见上图,FMEN的峰值内存占用位于HFAN内部:包含全局残差特征、输入特征、注意力以及输出特征,约4 times C times H times W

Enhanced Residual Block

自EDSR之外,ResBlock(见下图a)成为SISR中最常用模块。然而,ResBlock中的跳过连接会引入额外的内存占用并降低推理速度(实验表明:移除EDSR中的跳过连接,推理耗时将减少10%)。

为继承残差学习的优点且避免上述问题,作者设计了一种ERB模块(见上图b):它由两个RRRB与一个LReLU构成。RRRB受RepVGG与RCAN而来,在推理时可以折叠一个单个卷积。

High-frequency Attention Block

注意力机制已在SR领域得到广泛研究,可划分为Channel Attention、Spatial Attention、Pixel Attention以及Layer Attention。考虑到额外的内存占用与推理速度问题,我们设计了上图所示的HFAB:首先采用3times 3 卷积进行降维,然后采用ERB捕获局部县惯性,其次采用3times 3 卷积进行升维,最后sigmoid激活并与输入相乘。

注:在attention分支,作者引入了BN层。尽管已有研究(如EDSR)表明:BN会导致伪影问题。自此之后,BN不适用于SR已成为了常识。BN真的不能用于SISR吗?年前曾有一篇paper对此进行了思考,可参考:AdaDM: 让超分网络不仅可以用BN,性能还可以更优

回归正题,卷积的感受野比较有限,而在attention分支引入BN不仅有益于注入全局交互信息,同时有益于sigmoid的非饱和区域。在推理阶段,BN可以与后接或前接卷积合并。因此,在推理阶段,HFAB仅包含四个高度优化的算子:3times 3 卷积、LReLU、Sigmoid以及点乘。

2Experiments

上图对比了基于ResBlock、PlainBlock以及ERB的EDSR与FMEN性能,可以看到:相比ResBlock,ERB具有与其相当的性能,而PlainBloc则出现了严重性能下降。而在推理阶段,ERB可以折叠而PlainBlock形式并享受其高推理效率特征

上表对比了ResBlock与ERB在不同超分模型的推理效率对比,可以看到:相比ResBlock,ERB模块的推理效率平均快10%左右

上图从频域角度对HFAB进行了分析,可以看到:

  • 经HFAB处理后,输入特征的低频信号被抑制、高频信号被增强;
  • HFAB的输入会缓慢朝高频信号聚焦,而HFAB可以临时校正感兴趣区域;
  • 频域分离特征可以通过后续HFAB模块进一步增强。

上表对IMDN的CCA、RFDN的ESA以及本文的Attention进行了对比,可以看到:

  • 相比Baseline,无论CCA还是ESA均可以提升模型性能,但均低于本文的Attention方案;
  • 此外,由于多分支结构与低效算子导致CCA与ESA需要更多的计算耗时,而本文所设计Attention机制对于模型推理更友好,耗时几乎可忽略不计。

上表&图给出了不同方案的性能与效果对比,可以看到:FMEN取得了与LatticeNet相当的性能,同时具有更快推理速度(46ms vs 68ms)、更低内存占用(68M vs 225M)

上表给出了NTIRE2022EISR竞赛中不同方案的性能对比,可以看到:

  • 相比AIM2020冠军方案E-RFDN,FMEN-S参数量减少21.2%、FLOPs减少17.8%、推理耗时减少33.1%、峰值内存占用减少74%、Activation减少35.7%,而PSNR仅降低0.04dB
  • 相比NTIRE2022竞赛其他方案,FMEN-S取得了内存占用最低、Activation最低、卷积数最少、推理耗时第二的性能

3点滴反思

看完paper与code后,有这么几点反思分享给大家:

  • 以往方案的设计理念:attention部分使用尽可能少的参数量与计算量,尽可能提升主分支的计算量与参数量。而FMEN反其道而行:加大了attention部分的卷积数,减少主干部分卷积的数量。这是不是意味着Attention分支的作用被其“精简结构”限制了?进一步加大其结构到与主干部分相当的水平是否能够进一步提升性能呢?Attention与主干的相互作用又是怎样的?如果移除了Attention分支的Sigmoid,那么何为主、何为attention呢?
  • BN用于Attention竟然可以提升性能!已有SISR中的attention部分往往是将从架构设计那里“借鉴”并移除BN后使用。但是,FMEN与AdaDM则反其道而行之,将BN又一次进行了“引进”并提升性能。
  • 无论是BN-Conv合并还是Conv-BN合并,大家应该都不陌生。当然,大家更熟悉的是Conv-BN这种顺序关系,BN-Conv反而用的更少一些。一直认为:无论BN-Conv还是Conv-BN折叠前后是否等价与padding无关。被FMEN的实现打脸了,要想保证BN-Conv折叠前后的等价性,在折叠前需要在BN-Conv中间插入一个padding,且padding的数值与BN的参数信息有关

0 人点赞