极市导读
图像复原领域又一力作 >>加入极市CV技术交流群,走在计算机视觉的最前沿
arXiv:https://arXiv.org/abs/2111.09881
code:https://github.com/swz30/Restormer
本文是MPRNet与MIRNet的作者在图像复原领域的又一力作,也是Transformer技术在low-level领域的又一个SOTA。针对Transformer在高分辨率图像复原中存在的难点,提出了两种MDTA与GDFN两种改进,极大程度上缓解了计算量与GPU缓存占用问题。所提方案刷新了多个图像复原任务的SOTA性能。
通过MHSA与FFN进行改进,本文提出一种高效Transformer,它可以捕获长距离像素相关性,同时可适用于大尺寸图像。所提方案Restormer(Restoration Transformer)在多个图像复原任务上取得了SOTA性能,包含图像去雨、图像去运动模糊、图像去散焦模糊以及图像降噪(包含合成与真实噪声),可参见下图。
本文主要贡献包含以下几点:
本文旨在设计一种高效Transformer模型,它可以处理复原任务中的高分辨率图像。为缓解计算瓶颈,我们对MHSA进行了关键性改进并引入多尺度分层模块,见下图。
Overall Pipeline 给定退化图像 , Restormer的处理流程如下:
Transformer的主要计算负担来自自注意力层。在常规自注意力中,key-query的点乘操作会随分辨率以 增长。因此,基本不可能将自注意力层用于图像复原任务(图像复原的输入分辨率一般都比较大)。
为解决该问题,我们提出了MDTA(见上图),它具有线性复杂度,其关键成分在于:在通道维度(而非空间维度)执行自注意力计算跨通道的交叉协方差以生成关于全局上下文的隐式注意力特征图。作为MDTA的另一个重要成分,在计算特征协方差生成全局特征图之前,我们引入了深度卷积以突出局部上下文。
对于输入 , MDTA的处理流程如下:
MDTA的计算公式可以表示为:
注: 是一个可学习缩放参数,用于控制点积的幅值。
为进行特征变换, 常规FFN(由两个 卷积构成)对每个像素进行独立处理。本文则对齐进行了两个改动以 提升表达能力(见上图):
总而言之, GDFN可以控制信息流动, 进而使得每层聚焦于不同的细节信息。由于GDFN具有比FFN更多的 计算量, 我们降低扩展因子火以控制参数量与计算量。
基于CNN的复原模型通过采用固定尺寸图像块进行训练。然而,Transformer模型在较小块上训练可能无法进行全局统计信息编码,进而导致全分辨率测试时的次优性能。
针对该问题,我们提出了Progressive Learning机制:在训练的初期,模型在较小图像块上进行训练;在训练的后期,模型采用更大图像块进行训练 。由于更大的图像块会导致更长的计算耗时,我们随图像块提升降低batch参数以保持与固定块训练相当的耗时。
通过混合尺寸图像块训练的模型具有更优的性能。Progressive学习策略具有类似Curriculum学习策略相似的行为。
我们在不同的任务上进行了所提方案的性能验证,包含图像去雨、图像去运动模糊、图像去散焦模糊、图像降噪。
上图&表给出了所提方案在去雨任务上的性能与效果对比,可以看到:
上表给出了所提方案在不同去模糊数据集上的性能对比,可以看到:
上图&表给出了去散焦模糊任务上的性能对比,从中可以看到:
上表&图给出了不同图像降噪任务上的性能与效果对比,从中可以看到:
上表给出了关于模块的消融实验,可以看到:
从上表可以看到:
公众号后台回复“数据集”获取50+深度学习数据集下载~
# CV技术社群邀请函 #
备注:姓名-学校/公司-研究方向-城市(如:小极-北大-目标检测-深圳)
即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群
每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~