电话咨询

电话咨询:010-62607600

微信咨询

微信扫码

图像恢复 | Uformer:一种用于图像恢复的通用U形变压器

行业资讯 发布时间:2022-01-11 16:19:30 作者:深延科技

2021年6月6日提交到arxiv上的文章。ICCV2021的Eformer就是Uformer基础上改进而来,看来还是值得一读,简单记录下。

架构和创新点

架构如图:

image.png

比起普通的UNet,不同之处在于采用了LeWin Transformer,这种Transformer也是本工作的创新点。

所谓LeWin Transformer,就是local-enhanced window Transformer,其中包含W-MSA和LeFF:
W-MSA:non-overlapping window-based self-attention,作用是减小计算开销(传统transformer是在全局计算self-attention,而它不是);
LeFF:传统transformer中采用前馈神经网络,不能很好利用local context,LeFF的采用可以capture local information。

⚠️两个创新点:
提出LeWin Transformer,引入UNet
三种跳跃连接

主要模块细节

1、W-MSA

(本以为这是Uformer发明的,经提醒,swin Transformer里就有)

image.png

首先将将C×H×W的X分为N个C×M×M个patch,每个patch视为有M×M个C维vector(N = H × W / M²),这C个vector就输入W-MSA中。根据上述公式,简单理解就是将X分为不重叠的N片,然后对每一片进行self-attention的计算即可。

image.png

作者表示,虽说是在一片上进行self-attention的计算,但是在UNet的encode阶段,由于下采样的存在,所以在这一片上计算自注意力,对应在下采样前更大感受野上计算自注意力。

采用了relative position encoding,所以计算公式可以表示为:

image.png

这种位置编码的引用[48,41]分别是:
[48] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. arXiv preprint arXiv:1803.02155, 2018.
[41] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103.14030, 2021.

2、LeFF

LeFF是Incorporating Convolution Designs into Visual Transformers发明的,其中的Convolution-enhanced image Transformer (CeiT)包含了这个设计。

image.png

精髓在于对self-attention计算模块输出的N个token(vector),重新排列为 图片的“image”,然后进行depth-wise的卷积操作。看完CeiT作者给出的图解,再看Uformer作者给出的图解,就不难理解含义了:

image.png

每个线性层/卷积层之后,用的都是GELU激活函数。(depth-wise的卷积网上一搜就有,作用是减少参数,提升计算速度)

3、三种跳跃连接

UNet架构是有跳跃连接的,在本工作中就是将encoder部分的Transformer的输出传递到decoder部分,但是利用这些跳跃连接传递的信息的方式有很多种,作者探索了三种:

image.png

image.png

作者认为三种差不多,但第一种好一点点,所以采用第一种作为Uformer的默认设置。
Uformer的主要创新点就这些,其他内容不细看了。

计算开销

既然LeWin Transformer中的W-MSA主打减小计算开销,那么自然要看下算法复杂度:

给定feature map X,维度为C×H×W,如果是传统的self-attention,那么复杂度为 图片 ,分割为M×M的patch再做self-attention,则是 图片 ,复杂度有所减少。

实验结果

作者做了去噪、去雨、去模糊的实验。

image.png
image.png
image.png

来源:知乎
作者:umbrellalalalalaimage.png

立即使用深延AI平台