Abstract
轻量化与高效性是推动图像超分辨率(SR)算法落地应用的核心动力。本文提出一种简洁有效的轻量化图像超分辨率方案 ShuffleMixer,该方法利用大核卷积与通道拆分 - 重排算子开展特征学习。
现有超分模型大多通过堆叠多个小核卷积或者复杂算子来学习特征表示,与之不同,本文面向移动端友好的超分任务,探索基于大核卷积网络的设计思路。具体来说,本文设计由大核深度卷积以及两层基于通道拆分、重排的投影层构成基础模块,实现高效的特征混合。
由于自然图像内容具有强局部相关性,仅依靠大核深度卷积难以重建精细纹理细节。为解决该问题,同时保留模块的高效特性,本文在网络中引入融合倒残差卷积(Fused-MBConv),对不同特征之间的局部关联进行建模。
实验结果表明:所提 ShuffleMixer 在取得具有竞争力重建效果的同时,参数量与浮点运算量(FLOPs)相比现有前沿方法降低约 6 倍。在 NTIRE 2022 高效图像超分辨率挑战赛中,该主方案斩获模型复杂度赛道冠军 [23]。 代码开源地址:GitHub - sunny2109/MobileSR-NTIRE2022: [NTIRE 2022, EfficientSR] MobileSR: A Mobile-friendly Transformer for Efficient Image Super-Resolution · GitHub[NTIRE 2022, EfficientSR] MobileSR: A Mobile-friendly Transformer for Efficient Image Super-Resolution - sunny2109/MobileSR-NTIRE2022https://github.com/sunny2109/MobileSR-NTIRE2022
1 Introduction
单图像超分辨率(SISR)旨在从一张低分辨率图像重建得到对应的高分辨率图像。随着超高清电视、三星 Galaxy S22 Ultra、iPhone 13 Pro Max、华为 P50 Pro 等高清设备快速普及,该经典问题受到广泛关注。因此,研究高效且性能优异的算法来重建高分辨率图像,以适配这类设备的显示需求,具有重要的研究价值。
近年来,基于卷积神经网络(CNN)的超分辨率模型 [8, 9, 1, 16, 25, 45] 取得了优异的重建效果。但这类网络采用分层方式提取局部特征,往往需要堆叠更深、结构更复杂的模型来扩大感受野,以此提升重建性能。高昂的计算开销,导致这类重型超分模型难以部署在算力资源受限的移动端设备上,限制了实际落地应用 [44]。
为缓解超分模型参数量大、计算繁重的问题,研究人员已经提出多种方案以降低模型复杂度、提升推理速度,包括高效算子设计 [32, 28, 36, 9, 16, 1, 33, 43, 23, 27]、神经网络架构搜索 [6, 35]、知识蒸馏 [12, 13] 以及结构重参数化方法 [7, 23, 44]。这类方法大多基于改进小型空间卷积或者采用先进训练策略,对大核卷积的探索较少。此外,现有方法通常只侧重单一效率指标,在真实资源受限场景下表现欠佳。因此,亟需在模型复杂度、推理时延与超分重建质量三者之间实现更优的权衡。
图 1:在 B100 数据集 [2] 的 ×4 超分任务上,本文所提 ShuffleMixer 系列模型与其他轻量化方法的模型复杂度及重建性能对比。圆圈大小代表参数量,ShuffleMixer 实现了更优的权衡。
大感受野能够容纳更多的特征交互,对于超分辨率这类需要逐像素密集预测的任务而言,有助于重建出更加精细的结果。近期基于视觉 Transformer(ViT)的方法 [10,26,30,24] 采用多头自注意力(MHSA)机制学习长距离特征表征,在各类视觉任务上取得了前沿性能。然而,多头自注意力并不适合轻量化超分网络用来扩展感受野:其计算复杂度随输入分辨率呈二次增长(超分训练过程中输入图像分辨率通常较大且固定)。 采用大核标准卷积是另一种获取大感受野的简单方案,但模型开销较高。为让大核卷积能够落地应用,大核深度卷积[27,39,7] 是一种有效的替代方案。深度卷积在空间位置上共享权重、各通道之间相互独立,该特性也导致它难以捕捉充分的通道交互。因此,在轻量化网络设计中,提升深度卷积(DW Conv)的特征学习能力十分关键。
本文提出一种简洁高效的网络 ShuffleMixer,将大核卷积引入轻量化超分辨率网络设计。其核心思路是:在一个特征混合块内融合非局部与局部空间信息,同时控制参数量与浮点运算量。具体而言,本文采用大核深度卷积,聚合大范围区域内的空间信息;在通道混合方面,引入通道拆分与重排策略,降低参数量和计算开销,同时提升网络表征能力。基于上述两种算子,本文构建高效的 Shuffle Mixer 层。为进一步提升模型学习能力,在混合层中嵌入融合倒残差卷积(Fused-MBConv),强化特征的局部关联建模。综上,ShuffleMixer 依靠结构简洁的模块即可取得前沿重建性能。如图 1 所示,在现有各类轻量化超分方法中,ShuffleMixer 以最低的参数量与计算量实现了更优的性能权衡。
本文的主要贡献总结如下: (1) 探索大核卷积网络,挖掘图像超分任务中更多有效信息,构建了一种高效的超分辨率网络; (2) 提出通道拆分与重排算子,实现通道投影维度上的高效特征混合; (3) 为充分挖掘 Shuffle Mixer 层中跨组特征之间的局部关联,在超分网络中引入融合倒残差卷积(Fused-MBConv)。 本文将上述模块整合为一个端到端可训练网络,命名为 ShuffleMixer。实验结果表明,ShuffleMixer 的参数量与浮点运算量仅为前沿方法的约 1/6,同时可以取得和前沿方法相当的重建性能。
2 Related Work
基于卷积神经网络的高效超分方法基于 CNN 的各类方法采用不同思路降低模型复杂度。FSRCNN [9] 与 ESPCN [33] 采用后置上采样层,大幅降低固定输入尺寸带来的计算开销。Namhyuk 等人 [1] 在递归网络基础上结合分组卷积与级联连接,以节约参数量。Hui 等人 [16] 提出轻量化信息多蒸馏网络(IMDN),依靠特征拆分与拼接实现特征聚合;其改进版本 [43,23] 分别夺得 AIM2020 与 NTIRE2022 高效超分赛道冠军。与此同时,神经网络架构搜索也是热门方向:将网络结构搜索建模为带约束的多目标优化进化问题 [6,35]。另一类思路是借助知识蒸馏对重型深度模型做压缩与加速 [13,12]。值得注意的是,参数量更少、FLOPs 更低并不等同于移动端推理速度更快。因为浮点运算量指标忽略了诸多影响时延的关键因素,例如访存开销(MAC)与并行度 [28,32]。本文分析了影响超分模型实际推理效率的各类因素,设计了面向移动端部署的超分网络。
基于 Transformer 的超分方法Transformer 最早面向自然语言任务提出,通过堆叠多头自注意力与前馈 MLP 层,学习输入之间的长距离依赖关系。Dosovitskiy 等人 [10] 首次将视觉 Transformer 应用于图像识别任务。自此,基于 ViT 的模型在高层与底层视觉任务中得到越来越广泛的使用。3 Proposed Method在图像超分辨率方向,Chen 等人 [4] 提出预训练图像处理 Transformer(IPT),将原始 ViT 直接作用于无重叠分块嵌入特征。Liang 等人 [24] 借鉴 Swin Transformer [26] 的思路,提出基于窗口自注意力的图像复原模型,并取得优异效果。窗口自注意力相比全局自注意力计算效率更高,但仍然存在推理耗时久、显存开销大的问题。
大核卷积模型AlexNet [20] 是经典的大核卷积网络,启发了后续大量研究。全局卷积网络 GCN [31] 采用对称可分离大卷积核提升语义分割性能。但由于计算开销大、参数量多,自 VGG-Net [34] 之后,大尺寸卷积核便不再流行。 而随着高效卷积技术、Transformer 与 MLP 等新架构的发展,大核卷积重新受到学界关注。ConvMixer [39] 使用大核深度卷积替换 ViT [26,10] 或 MLP [38] 中的混合模块;ConvNeXt [27] 采用 7×7 深度卷积核对标准 ResNet 进行改造,取得和 Transformer 相当的效果;RepLKNet [7] 将卷积核扩大至 31×31,搭建纯卷积网络,在 ImageNet 上效果优于 Swin Transformer [26]。 不同于上述面向高层视觉任务、侧重构建大模型的方法,本文探索将大卷积核用于图像超分辨率轻量化网络设计的可行性。
3 Proposed Method
本文旨在面向 SISR 任务构建一款高效的大核卷积网络。为满足轻量化的设计目标,本文在特征混合模块中引入多项核心设计,实现信息的高效编码。本节首先详细介绍 ShuffleMixer 网络的整体流程;其次对作为网络基础组件的特征混合模块进行推导说明;最后给出训练损失函数的详细定义。
3.1 ShuffleMixer Architecture
ShuffleMixer 整体架构:给定低分辨率图像 \(I_{LR}\in \mathbb{R}^{C\times H\times W}\),其中 C、\(H\times W\) 分别代表通道数与空间分辨率;彩色图像的通道数 \(C=3\)。ShuffleMixer 首先通过一个 \(3\times3\)、输出通道数为 D 的卷积提取初始特征 \(Z_0\in \mathbb{R}^{D\times H\times W}\)。随后,本文设计特征混合模块(FMB),该模块包含两层 Shuffle Mixer 层以及 Fused-MBConv [36];模块以特征 \(Z_0\) 作为输入,输出更深层特征 \(Z_1\in \mathbb{R}^{D\times H\times W}\)。接着,由多个 FMB 串行输出的特征送入上采样模块,按照缩放因子 s 提升特征的空间分辨率。为最大限度减少放大模块的参数量,上采样单元仅由一层 \(1\times1\) 卷积与像素重排层(Pixel Shuffle)[33] 构成。对于 4 倍超分任务,通过两次 2 倍上采样逐级实现分辨率放大。最后利用卷积层将上采样后的特征映射为残差图像 \(I_R\in \mathbb{R}^{C\times sH\times sW}\);将该残差图与经过双线性插值放大的低分辨率图像相加,得到最终高分辨率图像 \(I_{SR}\):
\(I_{SR}=\uparrow_s(I_{LR})+I_R\) 式中 \(\uparrow_s(\cdot)\) 代表缩放因子为 s 的双线性插值操作。下文将对所提方法进行详细阐述。
图 2:本文所提 ShuffleMixer 的整体架构。该网络包含特征提取、特征混合与上采样三个部分,核心组件为特征混合块(FMB):每个 FMB 由两个 (a) Shuffle Mixer 层以及一个 (d) 融合倒残差卷积(Fused-MBConv,FMBConv)构成。每个 Shuffle Mixer 层包含两个 (b) 通道投影模块与一个大核深度卷积;通道投影模块内部集成通道拆分 - 重排操作、(c) 逐点 MLP 层、跳跃连接以及层归一化。
特征混合块用于挖掘局部与非局部信息,以获得更优质的特征表征。如图 2 (a) 所示,本文在特征 \(Z_0\) 上使用 Shuffle Mixer 层,高效实现非局部特征交互。对于每一个 Shuffle Mixer 层,采用大核深度卷积完成空间维度的特征混合。该算子能够以较少参数量获得较大的有效感受野,编码更加丰富的空间信息,从而重建完整、准确的图像结构。由表 3 的实验结果可见,在保证计算效率的前提下,采用更大尺寸的深度卷积核能够持续提升超分辨率重建性能。
为实现通道维度上的特征混合,本文采用逐点 MLP 完成通道投影。在深度卷积已确定的前提下,Shuffle Mixer 层的计算开销主要来自通道投影模块。为此,本文进一步引入通道拆分 - 重排(CSS)策略 [28] 提升该环节的效率。 具体来说,输入特征 \(Z_0\) 首先拆分为 \(Z_0^1\) 和 \(Z_0^2\);随后利用逐点 MLP 对拆分后的特征 \(Z_0^1\) 开展通道混合;最后通过通道重排操作,让拼接后的特征实现信息交互。该方案将通道投影层的参数量复杂度由 \(\Omega(4C^2)\) 降低至 \(\Omega(C^2)\)。该过程可以表述为:
式中,\(\sigma\) 为 SiLU 非线性激活函数 [11],\(W_0\)、\(W_1\) 代表逐点卷积;\(\text{Split}(\cdot)\) 和 \(\text{Shuffle}(\cdot)\) 分别表示在通道维度对特征进行拆分与重排。 该拆分操作存在一定局限:一半输入特征没有参与通道交互,模型表征能力受到制约;同时通道重排也无法保证所有特征都经过有效处理。受 MobileNetv2 模块 [32] 启发,本文将通道投影层重复设置,分别布置在大核深度卷积的前后两端,用于学习视觉特征表示。如表 2 消融实验所示,本文改进后的混合层取得与 ConvMixer 模块 [39] 相当的重建效果,同时消耗更少的参数量与浮点运算量。
自然图像内容具备局部相关性,堆叠的 FMB 模块无法充分挖掘局部特征,需要更大的模型容量来建模特征表示以提升超分性能。为此,本文在网络中嵌入卷积模块,强化特征的局部关联建模。具体而言,本文在每两层 Shuffle Mixer 层之后均匀插入融合倒残差卷积 Fused-MBConv。原始 Fused-MBConv 包含 3×3 扩张卷积、SE 层 [14](即常用的通道注意力)以及 1×1 压缩卷积。直接使用原始 Fused-MBConv 会带来可观的参数量与计算量开销,因此本文对其进行改造,适配轻量化需求。
本文首先移除 SE 模块,原因在于 SiLU 激活函数在一定程度上本身就具备门控机制。同时,隐层通道扩张会显著拖慢推理速度:原始扩张卷积通常采用很大的扩张系数(默认一般取 6),本文不再大幅度扩张通道,而是将输出通道数限制为 \(C+C'\)(实验中 \(C'\) 设置为 16),如图 2 (d) 所示。本文也对混合单元内的多种操作开展消融研究,详见 4.3 节。
3.2 Learning Strategy
为约束网络训练,最直接的思路是让网络输出图像的内容尽可能贴近真值图像:
式中,\(I_{SR}\) 和 \(I_{GT}\) 分别代表网络输出重建图像与对应的高分辨率真值图像。已有研究表明,仅依靠逐像素损失难以有效恢复高频纹理细节 [5]。因此本文引入频率约束对网络训练施加正则化。本文用于网络训练的损失函数定义如下:
式中,\(\mathcal{F}\) 代表快速傅里叶变换(Fast Fourier transform, FFT),\(\lambda\) 为权重系数,根据实验经验设置为 0.1。
4 Experimental Results
4.1 Datasets and implementation
数据集本文遵循已有工作 [22,24,23],在 DF2K 数据集上训练模型。DF2K 由 DIV2K [37] 与 Flickr2K [25] 合并而成,共包含 3450 张(800+2650)高质量图像。本文采用标准方案,对高分辨率真值图像进行双三次下采样,生成对应的低分辨率图像。测试阶段,在 Set5 [3]、Set14 [42]、B100 [2]、Urban100 [15]、Manga109 [29] 这五个公开基准数据集上,使用峰值信噪比(PSNR)与结构相似性(SSIM)指标评估模型性能。所有 PSNR 与 SSIM 均在 YCbCr 色彩空间的 Y 通道上计算。
实现细节本文在 RGB 通道下训练模型,并对输入图像块采用随机水平翻转与随机旋转进行数据增强。每个训练批次中,从低分辨率图像随机裁剪 64 张尺寸为 64×64 的图像块作为输入。模型以最小化 L1 损失与频域损失 [5] 为目标,采用 Adam 优化器 [19] 训练,总迭代次数为 300000 次。学习率固定设置为 \(5\times10^{-4}\)。所有实验基于 PyTorch 框架,在英伟达 Tesla V100 GPU 上完成。
本文依据特征通道数与深度卷积核尺寸设置两个版本模型,两种版本的 FMB 模块数量均为 5 个。ShuffleMixer 模型通道数为 64、卷积核大小 7×7;轻量化版本 ShuffleMixer-Tiny 通道数为 32、卷积核大小 3×3。训练代码与模型权重将会开源。
4.2 Comparisons with State-of-the-Art Methods
为验证本文方法的重建性能,将 ShuffleMixer 与多种前沿轻量化算法进行对比,包括 SRCNN [8]、FSRCNN [9]、VDSR [18]、DRCN [17]、LapSRN [21]、CARN [1]、EDSR-baseline [25]、FALSR-A [6]、IMDN [16]、LAPAR [22]、ECBSR [44] 以及 SMSR [40]。
表 1 给出了 ×2、×3、×4 放大倍率下,各模型在基准数据集上的定量对比结果。除 PSNR、SSIM 指标外,表中同时列出参数量与浮点运算量 FLOPs;FLOPs 测试条件为将图像超分至 1280×720 分辨率。 图 1 在 B100 数据集 ×4 超分任务上对比各模型的 FLOPs 与参数量。本文 ShuffleMixer 在参数量与计算量更低的条件下,取得具有竞争力的重建效果。特别地,ShuffleMixer 参数量与 CARN-M 接近,但在全部测试集上性能显著优于 CARN-M。轻量版本 ShuffleMixer-Tiny 仅 113K 参数量,性能仍优于诸多现有算法。对于 ×2 与 ×3 倍率任务,ShuffleMixer 系列同样可以取得相当的效果。
尽管 IMDN [16]、LAPAR-A [22] 和 SMSR [40] 可以取得相近的 PSNR/SSIM 指标,但 ShuffleMixer 的模型复杂度更低。本文同时在 ×4 超分任务上对比了 CARN [1]、CARN-M [1]、LAPAR-A [22] 等轻量化快速模型的 GPU 推理耗时,所提方法推理速度更快。ShuffleMixer-Tiny、ShuffleMixer 重建一张 1280×720 高分辨率图像分别仅需 0.016 s 和 0.021 s;作为对照,CARN-M、CARN、LAPAR-A 的推理耗时依次为 0.017 s、0.019 s、0.031 s。 值得注意的是,PyTorch 对大核深度卷积的原生支持较差;文献 [7] 指出,采用经过优化的深度卷积算子能够进一步加快本文模型的推理速度。上述结果充分验证了本文方法的有效性。
图 3:在 Set14 与 Urban100 数据集上的 ×4 超分视觉对比。定性结果表明,本文所提方法能够生成视觉效果更佳的重建图像,更好地恢复图像结构与细节。
本文进一步在真实低质量图像上验证所提方法。图 4 展示了历史图像数据集 [21] 中的一组样例。对比方法 [15,22] 的重建结果存在明显伪影;本文方法与 CARN [1] 均可得到平滑的细节,但本文重建图像的结构更加清晰。
4.3 Analysis and Discussions
ShuffleMixer 的核心在于 Shuffle Mixer 层、特征混合块以及大核卷积。本节基于 Tiny 版本模型,在 DIV2K 数据集 [37] 的 ×4 超分任务上,对上述组件分别开展消融验证。
特征混合块的有效性为验证特征混合块的有效性,本文以 CDC 模型作为基线。首先在每两个 Shuffle Mixer 层之后嵌入一层 3×3 卷积,相比基线模型 PSNR 提升 0.13 dB。为进一步探究不同特征融合方式的效果,本文对比 S-Conv 与 C-Conv 两种方案:S-Conv 对输入、输出特征逐元素相加,再接 3×3 卷积;C-Conv 在通道维度拼接输入、输出特征,再接 3×3 卷积。
由表 2 (b) 可见,两种方案均优于基线;C-Conv 取得更高 PSNR,但计算开销更大。图 5 为上采样模块之前沿通道维度取平均得到的特征图,该结果表明,增强特征元素之间的局部关联,有利于捕捉更精细的高频信息。 在 S-Conv 的基础上,本文继续将卷积层替换为基础残差块(S-ResBlock)与融合倒残差卷积(S-FMBConv)。如表 2 (b) 所示,S-FMBConv 在模型复杂度与超分性能之间取得均衡折中。因此本文选用 S-FMBConv 来增强特征之间的局部关联。
大核深度卷积的有效性为验证大卷积核带来的作用,本文选取 \(3\times3\) 至 \(13\times13\) 多种卷积核尺寸分别测试。表 3 结果表明,增大卷积核能够提升重建性能。其中,\(7\times7\) 深度卷积相比 \(3\times3\) 版本 PSNR 提升 0.07 dB,仅额外增加 12K 参数量与 0.8G FLOPs。同时可以观察到,当核尺寸超过 \(7\times7\) 后,性能增益变得十分微弱。因此,综合权衡重建精度与模型复杂度,本文最终采用 \(7\times7\) 的卷积核。
5 Conclusion
本文提出了一种轻量化深度网络用于图像超分辨率任务。该网络 ShuffleMixer 包含 Shuffle Mixer 层,依靠更大的有效感受野高效提取非局部特征表征;同时引入改造后的 Fused-MBConv,对 Shuffle Mixer 输出特征建模局部关联,该模块对提升超分性能至关重要。本文在通用基准数据集上对 ShuffleMixer 开展定性与定量测试。实验结果表明,相较于现有前沿方法,ShuffleMixer 在取得有竞争力重建效果的同时,具备更高的计算效率。
Broader Impact
本文是一项基于大核卷积网络、面向轻量化高效图像超分辨率方向的探索性研究。该方法可部署在资源受限的设备环境中用于图像质量提升,例如处理手机拍摄照片,或是在视频通话、会议场景下降低传输带宽。 但超分辨率技术也存在潜在负面影响:不法分子可能利用该技术增强人脸、人体特征,造成身份信息泄露。值得说明的是,图像超分技术带来的正向社会价值远大于其潜在风险。本文呼吁使用者在应用该技术及其衍生工具时,不得损害公众个人权益。