Transformer再下一城！low-level多个任务榜首被占领，中科大等联合提出：Uformer

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

转载自：AIWalker

今天要分享的是Low-level领域的第二发Tranformer，它直接占领了图像降噪、图像去雨、图像去模糊等领域的高峰。Transformer真可谓CV领域搅局者！

第一发工作，可见：CVPR 2021 | Transformer进军low-level视觉！北大华为等提出预训练模型IPT

标题&作者团队

Uformer: A General U-Shaped Transformer for Image Restoration

paper: https://arxiv.org/abs/2106.03106

code: https://github.com/ZhendongWang6/Uformer

Abstract

本文提出一种有效且高效的基于Transformer的架构Uformer用于图像复原，它采用Transformer模块构建了一种分层编解码网络。Uformer的两个核心设计使其适用于图像复原任务：

• 局部增强窗口Transformer模块：我们采用非重叠的基于窗口的自注意力降低计算量，同时在前向网络中采用深度卷积进一步改善其捕获局部上下文的能力；

• 跳过连接机制：我们对跳过连接机制进行了探索以有效的将编码器的信息传递到解码器。

受益于上述两种设计，Uformer具有为图像复原捕获有用依赖关系的能力。我们通过几个图像复原任务(包含图像降噪、图像去雨、图像去模糊、图像去摩尔纹)实验验证了所提Uformer的优异性能，比如：

• 在图像降噪方面，相比NBNet，Uformer在SIDD与DND基准数据集上分别取得了0.02dB与0.07dB指标提升，同时计算量减半；

• 在图像去雨方面，相比RCDNet，Uformer在SPAD数据及上取得了4.53dB指标提升，达到了46.00dB；

• 在去模糊、摩尔纹移除方面，所提Uformer分别取得了0.52dB与0.81dB指标提升。

Method

接下来，我们首先对所提Uformer的整体架构进行简单介绍；然后，我们对Uformer的基础模块LeWin Transformer模块进行详细介绍。最后，我们引入三种跳过连接变种进行编码器与解码器之间的信息桥接。

Overall Pipeline

上图给出了Uformer的整体架构示意图，它是一种带跳过连接的U型网络。具体来讲，给定输入退化图像 ，Uformer首先采用 卷积+LeakyReLU提取底层特征 ；接下来，按照U型架构，特征 将由K个编码器进行处理，每个编码器包含多个所提LeWin Transformer模块与一个下采样层。LeWin Transformer利用自注意力机制捕获长距离依赖，同时通过非重叠窗口上的自注意力降低计算量；在下采样层，我们首先将平展特征reshape为2D形式，然后采用stride=2的 卷积下采样，同时通道数加倍。比如，在给定输入 下，第 阶段的编码器输出特征图 。

然后，瓶颈阶段位于编码阶段的底部，它由多个LeWin Transformer模块构成。在这个阶段，受益于分层结构，Transformer可以捕获更长距离依赖，甚至全局依赖。

在特征重建方面，所提解码器同样包含K个阶段，每个阶段包含一个上采样层与多个LeWin Transformer模块(类似于编码器部分)。我们采用stride=2的 反卷积进行上采样，它在提升特征分辨率同时降低通道数。上采样之后的特征以及同级编码器的特征将送入到LeWin Transformer模块。经过K个解码器处理后，我们将平展特征reshape为2D特征，然后采用 卷积得到残差图像 。最后，重建图像通过 得到。损失函数选择了Charbonnier损失。

注：在实验中，我们设置K=4，每个阶段包含两个LeWin Transformer模块。

LeWin Transformer Block

采用Transformer进行图像复原有两个主要问题：

• 计算量：标准Transformer在所有词之间计算全局自注意力，这会导致非常大的计算量，对于高分辨率图像而言尤其不适用。

• 局部上下文信息：局部上下文信息对于图像复原非常重要，而已有研究表明：Transformer在捕获局部上下文信息方面比较弱。

为解决上述两个问题，我们提出了局部增强窗口Transformer(Locally-enhance Window, LeWin)，见前面的Figure1b，它采用Transformer的自注意力捕获长距离依赖，同时将卷积引入到Transformer以捕获有用的局部上下文信息。

具体来说，给定第 莫开的输出特征 ，我们基于以下两个核心设计构建模块：

• 非重叠窗口多头自注意力(W-MSA)

• 局部增强前馈网络(LeFF)

LeWin模块可以表示成如下形式：

W-MSA 不同于标准Transformer的全局自注意力，我们在非重叠局部窗口执行自注意力，这种方式可以有效降低计算量。给定2D输入特征 ，我们先将X拆分为大小为 的非重叠块并得到每个窗口的平展与转置特征 。接下来，我们在每个窗口特征上执行自注意力。第k个头计算过程可以描述如下：

所有头的结果进行拼接并经线性投影得到最后的结果。受启发于已有研究Swin，我们还在注意力模块中引入了相对位置编码，因此注意力计算过程可以表示成如下形式：

相比全局自注意力，基于窗口的自注意力可以大幅减少计算量，从 减少到 。正如下图a所示，由于Uformer采用了分层结构，在低分辨率特征时该注意力具有大感受野，可以有效捕获长距离依赖。

LeFF  正如已有研究指出：标准Tranformer中的前馈网络具有非常有限的局部上下文信息捕获能力。考虑近邻像素对于图像复原的重要性，我们在前馈网络中添加了一个深度卷积以捕获局部上下文信息，见上图b。我们首先采用线性投影层提升每个词的维度；然后我们将词reshape为2D特征并采用 卷积捕获局部信息；其次我们将上述特征平展并通过另一个线性层收缩维度以匹配输入通道。注：我们采用GELU激活函数，卷积后同样采用该激活函数。

Variants of Skip-Connection

为研究如何将编码部分的特征传递给解码部分，考虑到Transformer的自注意力计算，我们设计了三种不同的跳过连接机制：

• Concatenation-based Skip-Connection：基于拼接的跳过连接广泛应用于UNet架构中。为构建该网络，我们首先将第 编码阶段的平展特征 与第 解码阶段的特征 进行拼接；然后我们将拼接特征送入到第一个LeWin的W-MSA模块，见上面的Figure2-c1。

• Cross-based Skip-Connection：不同于直接拼接，我们参考语言Transformer中的解码结构设计了Cross-Skip，见上面Figure2-c2。我们首先在第一个LeWin模块中添加一个额外的注意力模块，它用于从解码特征 中寻找自相似像素信息，而第二个注意力则从解码特征 中提取key与values，从第一个注意力模块输出中提取queries。

• Concatenation-based Cross-attention as Skip-connection：对上述两种进行了组合设计了第三种跳过连接，见上面的Figure2-c3.

Experiments

接下来，我们先对所提Uformer进行系统性研究，然后再通过不同任务验证所提方案的有效性。

Ablation Studies

在该部分，我们以SIDD数据集上的图像去噪为参考进行消融实验分解，结果见上表，从中可以看到：

• Transformer vs Convolution 当替换LeWin为ResBlock后得到了UNet架构。很明显，Uformer16取得了39.66dB指标，比UNet32高0.04dB；Uformer32取得了39.77dB指标，比UNet48高0.12dB。该结果验证了所提LeWin的有效性。

• Hierarchical Structure vs Single Scale 对比单尺度的ViT与本文的分层架构Uformer，Uformer具有更高的指标，更低的参数量。该结果验证了分层特征对于图像复原的重要性。

• Which Skip-connection is better 三种跳过连接方式取得相似的指标，其中Concatenation方式取得稍高性能。我们默认采用Concatenation方式。

• Where to enhance locality 相比baseline，在前馈网络中引入局部性可以带来0.03dB指标提升，而在自注意力部分引入局部性会造成0.02dB下降。这意味着前馈网络更适合引入局部性。

• Is shifted window important 窗口移位可以带来0.01dB指标提升。

Real Noise Removal

上表给出了SIDD与DND数据上的图像降噪性能对比，从中可以看到：

• 在SIDD数据方面，Uformer取得了最佳的39.77dB指标，超过了之前最佳NBNet的39.75dB；

• 在DND数据方面，Uformer以0.07dB指标超过了之前最佳方案NBNet。

• 相比其他方案，Uformer具有更低的计算复杂度，更高的指标。

上图给出了SIDD与DND数据上的视觉效果对比，可以看到：Uformer不仅成功的移除了噪声，同时很好的保留了纹理细节。

Real Rain Removal

上表给出了SPAD数据上的图像去雨性能对比，从中可以看到：相比现有工作，Uformer取得了近5dB指标提升。视觉效果见下图，Uformer可以成功的移除雨点同时引入最少的伪影。

Defocus Blur & Moire Pattern Removal

上图给出了虚焦模糊与摩尔纹移除性能与视觉效果对比，从中可以看到：

• 在虚焦模糊移除方面，所提Uformer取得了最佳指标25.65dB，同时具有更佳的视觉效果；

• 在摩尔纹移除方面，所提Uformer取得了最佳指标27.30dB，同时具有最佳的视觉效果、更少的伪影。

Uformer论文PDF下载

后台回复：Uformer，即可下载论文PDF

CVPR和Transformer资料下载

后台回复：CVPR2021，即可下载CVPR 2021论文和代码开源的论文合集
后台回复：Transformer综述，即可下载最新的两篇Transformer综述PDF
CVer-Transformer交流群成立
扫码添加CVer助手，可申请加入CVer-Transformer 微信交流群，方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch和TensorFlow等群。
一定要备注：研究方向+地点+学校/公司+昵称（如Transformer+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群
▲长按加小助手微信，进交流群▲点击上方卡片，关注CVer公众号
整理不易，请点赞和在看