arXiv:https://arXiv.org/abs/2111.09881
code:https://github.com/swz30/Restormer
本文是MPRNet与MIRNet的作者在图像复原领域的又一力作,也是Transformer技术在low-level领域的又一个SOTA。针对Transformer在高分辨率图像复原中存在的难点,提出了两种MDTA与GDFN两种改进,极大程度上缓解了计算量与GPU缓存占用问题。所提方案刷新了多个图像复原任务的SOTA性能。
Abstract
通过MHSA与FFN进行改进,本文提出一种高效Transformer,它可以捕获长距离像素相关性,同时可适用于大尺寸图像。所提方案Restormer(Restoration Transformer)在多个图像复原任务上取得了SOTA性能,包含图像去雨、图像去运动模糊、图像去散焦模糊以及图像降噪(包含合成与真实噪声),可参见下图。
本文主要贡献包含以下几点:
- 提出了一种编解码Transformer用于高分辨率图像上多尺度local-global表达学习,且无需进行局部窗口拆分;
- 提出一种MDTA(Multi-Dconv head Transposed Attention)模块,它有助于进行局部与非局部相关像素聚合,可以高效的进行高分辨率图像处理;
- 提出一种GDFN(Gated-Dconv Feed-forward Network)模块,它可以执行可控特征变换,即抑制低信息特征,仅保留有用信息。
Method
本文旨在设计一种高效Transformer模型,它可以处理复原任务中的高分辨率图像。为缓解计算瓶颈,我们对MHSA进行了关键性改进并引入多尺度分层模块,见下图。
Multi-Dconv Head Transposed Attention
Transformer的主要计算负担来自自注意力层。在常规自注意力中,key-query的点乘操作会随分辨率以增长。因此,基本不可能将自注意力层用于图像复原任务(图像复原的输入分辨率一般都比较大)。
为解决该问题,我们提出了MDTA(见上图),它具有线性复杂度,其关键成分在于:在通道维度(而非空间维度)执行自注意力计算跨通道的交叉协方差以生成关于全局上下文的隐式注意力特征图。
作为MDTA的另一个重要成分,在计算特征协方差生成全局特征图之前,我们引入了深度卷积以突出局部上下文。
Gated-Dconv Feed-forward Network
Progressive Learning
基于CNN的复原模型通过采用固定尺寸图像块进行训练。然而,Transformer模型在较小块上训练可能无法进行全局统计信息编码,进而导致全分辨率测试时的次优性能。
针对该问题,我们提出了Progressive Learning机制:在训练的初期,模型在较小图像块上进行训练;在训练的后期,模型采用更大图像块进行训练。由于更大的图像块会导致更长的计算耗时,我们随图像块提升降低batch参数以保持与固定块训练相当的耗时。
通过混合尺寸图像块训练的模型具有更优的性能。Progressive学习策略具有类似Curriculum学习策略相似的行为。
Experiments
我们在不同的任务上进行了所提方案的性能验证,包含图像去雨、图像去运动模糊、图像去散焦模糊、图像降噪。
Image Deraining
上图&表给出了所提方案在去雨任务上的性能与效果对比,可以看到:
- 相比此前最佳SPAIR,Restormer在所有数据集上取得了平均1.05dB指标提升;
- 在Rain100L数据集上,性能增益甚至高达2.06dB;
- Restormer可以生成更好的无雨图像,且可以有效的保持结构内容。
Motion Deblurring
上表给出了所提方案在不同去模糊数据集上的性能对比,可以看到:
- 相比MIMO-UNet+,所提Restormer可以取得了平均0.47dB指标提升;
- 相比MPRNet,所提方案Restormer可以取得平均0.26dB指标提升;Restormer的FLOPs仅为MPRNet的81%;
- 相比IPT,所提方案Restormer取得了0.4dB指标提升,同时具有更少的参数量(4.4x)、更快的推理速度(29x);
- 可视化效果见下图,很明显:Restormer重建结果更锐利清晰。
Defocus Deblurring
上图&表给出了去散焦模糊任务上的性能对比,从中可以看到:
- 无论是单帧还是双摄图像,所提方案大幅均优于其他方案;
- 在组合场景方面,相比此前最佳IFAN,所提方案取得了约0.6dB指标提升;相比Uformer,所提方案取得了1.01dB指标提升;
- 所提方案可以有效移除空间可变的散焦模糊。
Image Denoising
上表&图给出了不同图像降噪任务上的性能与效果对比,从中可以看到:
- Gaussian Denoising:所提方案在两种实验配置下均取得了SOTA性能。对于极具挑战性的50噪声水平的Urban100数据,Restormer取得了比DRUNet高0.37dB指标,比SwinIR高0.31dB。此外,相比SwinIR,Restormer计算量更少,速度更快。
- Real Denoising:所提方案是仅有的指标超过40dB的方案。相比此前最佳MIRNet与Uformer,所提方案分贝取得了0.3dB与0.25dB指标提升。
- Visual:从视觉效果上来看,所提方案可以重建更清晰的图像,同时具有更细粒度的纹理。
Ablation Study
上表给出了关于模块的消融实验,可以看到:
- 相比基线模型,MDTA可以带来0.32dB指标提升
- GDFN可以在MTA基础上取得0.26dB指标提升;
- GDFN与MDTA的组合可以取得了0.51dB指标提升。
从上表可以看到:
- 在提炼阶段添加Transformer可以进一步提升模型性能;
- Progressive学习机制可以取得更佳的指标,提升约0.07dB;
- 深而窄的模型比宽而浅的模型质保更高。
来源:AIWalker
作者: HappyAIWalker
推荐阅读
- DAMO-YOLO | 超越所有YOLO,魔搭社区开源至强YOLO,5行代码即可体验!
- 超简Transformer风格ConvNet,比ConvNeXt还秀,媲美VAN
- 端侧高效骨干RepGhost | 重参数赋能GhostNet,达成端侧超高效骨干
- PP-YOLOE-R | 手把手教你将PP-YOLOE用于旋转目标检测,顺带达成SOTA性能!
本文章著作权归作者所有,任何形式的转载都请注明出处。更多动态滤波,图像质量,超分辨相关请关注我的专栏AIWalker。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。