Restormer: Efficient Transformer for High-Resolution Image Restoration

1. Motivation and Background

1.1 CNNs及其局限性

  图像修复在于去除低质量图片上的噪音以重建高质量的图片,而修复图像需要强大的图像先验。CNNs在学习广尺度数据的泛化先验上有着良好的性能,故近年出现了很多基于CNNs的方法,和传统修复方法相较,人们更偏爱于前者。“卷积”是CNN基本操作,它能够提供局部连接和平移等方差,这使得其带来在效率和泛化能力上提升的同时,也导致了两个问题:卷积操作的感受野有限,这使得它不能建立大范围的像素关系和卷积过滤的权重在理论推导时是静态的,因此不能随着输入内容的变化而变化。为了解决这个问题,动态、强大的自注意力机制(SA)是一个很好的选择,它能够通过给输入像素赋予不同权重再求和来计算对应输出。

1.2 自注意力机制(SA:self-attention)

  自注意力机制是Transformer的核心组件,它的实现和CNN不同。例如,多头注意力对并行化和高效表征学习进行了优化。Transformer在自然语言处理、高层视觉等任务上有着state-of-the-art的性能。尽管SA在捕获广域像素交互上性能较好,但是它的复杂度会随着空间分辨率的提高呈平方式增长,故直接将其用于高分辨率图像处理(通常是图像修复)上是不可行的。

  所以虽然Transformer模型减轻了cnn的缺点(即有限的接受域和对输入内容的不适应性),但其计算复杂度随着空间分辨率呈二次增长,因此无法应用于大多数涉及高分辨率图像的图像恢复任务。将transformer调整应用于图像修复上的工作,为了减小计算负荷,这些方法要么将SA用于每个像素周围8x8的小空间窗口,要么将输入图像分割成无重叠的48x48的块,再在单个块上计算SA。但是这对SA空间内容的限制和捕获真正的广域像素关系是相违背的,特别是对于高分辨率图像来说。

1.3 Restormer

  为解决上述问题,本论文提出了一个有效的Transformer模型,通过在构建块(多头注意和前馈网络)中进行几个关键设计,使其可以捕获远程像素交互,同时仍然适用于大图像。模型会在几个图像恢复任务上实现了最先进的结果,包括图像去噪、单图像运动去模糊、离焦去模糊(单图像和双像素数据)和图像去噪(高斯灰度/彩色去噪和真实图像去噪)。

2. Method

2.1 Overall Pipline

  Restormer的整体执行流程如下:首先输入图像大小为 I∈R^{H×W×3} ,利用一个卷积操作获得特征嵌入Fo∈R^{H×W×C} 。接着F0通过一个对称的4层编码-解码结构,得到高维特征Fd∈RH×W×2C,每一层编码/解码都包括多个Transformer 模块,从上到下,每一层中的Transformer模块数量逐渐递增,分辨率逐渐递减。其次在编码-解码器之间使用跳跃连接来传递低维特征信息。Fd进一步经过Refinement模块来提取细节特征。最后经过一个卷积层,并与输入图像进行叠加,得到最后的输出图像。

  Restormer的整体结构如下图,包括四层对称的编码器-解码器。其核心组件包括Mutli-Dconv head Transposed Attention(MDTA)和Gated-Dconv Feed-forward Network(GDFN)

Restormer的整体结构

2.2 Mutli-Dconv head Transposed Attention

  为了缓解传统SA中点积交互的时间和内存复杂度与输入空间分辨率成二次方增长,即对于W×H像素的图像复杂度为O(W2×H2) ,本文通过把HW×WH的attention优化成通道×通道的attention,使得MDTA具有线性复杂度,它有助于进行局部和非局部相关像素聚合,实现高效的高分辨率图像处理。Restormer模型在获得归一化张量Y后,通过1×1卷积聚合跨通道的上下文,再使用3×3深度卷积来编码通道级上下文,生成了query、key、value(即图中的\hat{Q}, \hat{K}, \hat{V})投影。接着应用softmax函数控制K和Q的点积大小,并通过SA计算生成转置注意图A。

  MDTA的过程被定义为如下公式:

MDTA公式

  该模块结构示意图和过程定义如下:

MDTA结构

2.3 Gated-Dconv Feed-forward Network

  在传统的前馈网络(FN)中,特征转换是通过在每个像素位置独立且相同地操作来实现的。它使用两个 1×1 卷积,一个用于扩展特征通道(通常扩大 γ=4 倍),另一个用于将通道数减少回原始输入维度。在隐藏层应用非线性激活函数。在这项研究中,我们提出了两个基本修改以改善特征表示学习:(1)门控机制和(2)深度可分离卷积。我们的 GDFN 架构如图2(b)所示。门控机制被构建为线性变换层的两个并行路径的逐元素乘积,其中一个路径使用 GELU 非线性激活函数激活。与MDTA类似,我们还在GDFN中包含了深度可分离卷积,以编码来自空间相邻像素位置的信息,这对于学习有效恢复所需的本地图像结构非常有用。

  给定输入张量X∈R^{\hat{H}×\hat{W}×\hat{C}},GDFN 的公式为:

GDFN公式

  其中 · 表示逐元素乘法,φ表示 GELU 非线性激活函数,LN 是层归一化操作。总体上,GDFN 控制着我们管道中各层次的信息流,从而使每个层次都能专注于与其他层次互补的精细细节。也就是说,与MDTA(专注于丰富上下文信息的特征)相比,GDFN 的作用更为独特。由于所提出的 GDFN 比常规 FN执行的操作更多,我们减小了扩张比例 γ,以保持相似的参数数量和计算负担。

2.4 Progressive Learning

  在裁剪后的小 patch 上训练 Transformer 模型可能不会对全局的图像统计信息进行编码,从而在测试时的全分辨率图像上提供次优的性能。因此本问使用了渐进式学习,其中网络在早期阶段用较小的图像 patch 上进行训练,在后期的训练阶段中逐渐增大。通过渐进式学习在混合大小的 patch 上训练的模型在测试时表现出增强的性能,其中图像可以具备不同的分辨率。渐进式学习策略的行为方式类似于课程学习过程:网络从简单的任务开始,逐渐转向学习更复杂的任务,同时保留精细的图像结构。由于对大 patch 的训练需要花费更长的时间,因此会随着 patch 大小的增加而减小批量大小,以保持每个优化步骤的时间与固定 patch 训练相似。

3. Experiment and Analysis

3.1 Experiment Setting

  Restormer 在四种图像处理任务上进行了实验评估:(a)图像去雨水、(b)单图像运动去模糊、(c)焦外模糊去除(针对单图像和双像素数据)、以及(d)图像去噪(针对合成和真实数据),同时为不同的图像恢复任务训练了单独的模型。

  我们从模型架构、精化阶段、训练参数、渐进式学习和数据增强五个方面说明Restormer的实验设置。Restormer 使用了一个包含四个级别的编码器-解码器结构。从第一级到第四级,Transformer 块的数量分别为 [4, 6, 6, 8],MDTA 中的注意力头数为 [1, 2, 4, 8],通道数量分别为 [48, 96, 192, 384]。在精化阶段,Restormer包含 4 个块,并设置GDFN 中的通道扩展因子:γ=2.66。实验使用了 AdamW 优化器(参数为 β1=0.9,β2=0.999,权重衰减为 1e−4)和 L1 损失函数进行了 300K 次迭代的训练,初始学习率为 3e−4,并逐渐减少到 1e−6(采用余弦退火 [51])。为了渐进式学习,实验从 128×128 的图块大小和批处理大小为 64 开始训练,随着迭代次数的增加,图块大小和批处理大小的配对逐步更新为 [(1602, 40), (1922, 32), (2562, 16), (3202, 8), (3842, 8)] 分别在 [92K, 156K, 204K, 240K, 276K] 次迭代时更新。最后使用水平和垂直翻转技术进行数据增强。

3.2 Experiment Result

3.2.1 Image Deraining

  5个去雨数据集上的图像去雨的结果如下表 1 和图 3 所示。

表1. 图像去雨的结果

图3. 图像去雨的例子

3.2.2 Single-image Motion Deblurring

  GoPro数据集上的单图像运动去模糊的结果如下表 2 所示和图 4 所示。

表2. 单图像运动去模糊的结果

图4. 单图像运动去模糊的例子

3.2.3 Defocus Deblurring

  DPDD测试集上的散焦去模糊的实验对比如下表 3 所示,包括单幅图像散焦去模糊和双像素散焦去模糊。双像素散焦去模糊对比案例如图 5 所示。

表3. DPDD测试集上的散焦去模糊的实验对比

图5. DPDD数据集上的双像素散焦去模糊比较

3.2.4 Image Denoising

  高斯灰度图像去噪在两类方法(CNN 变体 和 Transformer 变体)上的结果比较如下表 4 所示。

表4. 高斯灰度图像去噪在两类方法上的结果比较

  高斯彩色图像去噪的实验结果如下表 5 所示。

表5. 高斯彩色图像去噪的实验结果

  真实图像去噪的实验结果下表 6 所示。


表6. SIDD 和 DND 数据集上的真实图像去噪的实验结果

  图像去噪的视觉结果如下图 6 所示。

图6. 图像去噪的视觉结果

3.2.5 Ablation Studies

  Transformer 块对结果的消融实验如下表 7 所示。

表7. Transformer 块的消融实验

  在解码器层级1中的concat(w/o 1x1 conv)和细化阶段对实验结果的影响如下表 8 所示。

表8. 在编码器层级1中的concat和细化阶段对实验结果的影响

  在固定 patch 大小和逐渐增大 patch 大小上训练 Restormer 的结果对比如下表 9 所示。

表9 在固定 patch 大小和逐渐增大 patch 大小上训练 Restormer 的结果对比

  更深或更宽的模型对实验结果的影响如下表 10 所示。
表10. 更深或更宽模型的结果对比

4. Conclusion

  Restormer,它是计算效率高的处理高分辨率图像模型。我们介绍了Transformer模块核心组件的关键设计,以改进特征聚合和转换。具体来说,我们的multiDconv头部转置注意(MDTA)模块通过跨通道而不是空间维度应用自注意来隐式地模拟全局上下文,因此具有线性复杂性而不是二次复杂性。此外,所提出的门控-直流前馈网络(GDFN)引入了一种门控机制来进行受控特征转换。为了将cnn的强度整合到Transformer模型中,MDTA和GDFN模块都包括用于编码空间局部上下文的深度卷积。在16个基准数据集上的大量实验表明,Restormer在许多图像恢复任务中实现了最先进的性能。

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 205,236评论 6 478
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 87,867评论 2 381
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 151,715评论 0 340
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,899评论 1 278
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,895评论 5 368
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,733评论 1 283
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 38,085评论 3 399
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,722评论 0 258
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 43,025评论 1 300
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,696评论 2 323
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,816评论 1 333
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,447评论 4 322
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 39,057评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 30,009评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,254评论 1 260
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 45,204评论 2 352
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,561评论 2 343

推荐阅读更多精彩内容