Uformer: A General U-Shaped Transformer for Image Restoration
图像去噪、运动去模糊、散焦去模糊和去噪
CVPR2022
extensive experiments 大量的实验
capture long-range dependencies 获取长距离依赖
leverage the capability of - 利用- 的能力
To this end 为此
approach various perturbations 处理不同的扰动
calibrate
the features and encourage more details recovered:校准特征并鼓励恢复更多细节
comprehensive ablation studies 全面消融实验
摘要
提出基于Transformer架构的Uformer用于图像恢复的领域,使用Transformer build一个分层的编码器解码器网络。主要有两个创新:
1.introduce引入一个locally-enhanced window Transformer块,基于不相重合window窗口的自注意力机制而不是全局自注意力机制。(好处:当获取高分辨率图像的全局信息的时候,减少计算的复杂性。)
2.propose提出了多尺度空间偏置形式的可学习多尺度恢复调制器,以调整Uformer解码器的多个层中的特征。(优势:展示了在引入边际额外参数和计算成本的同时,为各种图像恢复任务恢复细节的卓越能力。)
在这两种设计的支持下,Uformer具有捕获图像恢复的局部和全局依赖性的强大能力。
introduction
动机:基于卷积ConvNets-based的sota工作虽然取得了不错的结果但是在获取长距离依赖信息时候有限制。现有的最近的工作在低分辨率特征图上应用自注意力层balabala低分辨率可以降低计算复杂度
这篇文章的解决方案:
**提出基于UNet架构的Ufromer。*对于全局来说,LeWin Transformer基于非重叠窗口的自注意力机制,而非全局注意力机制,显著降低高分辨率特征图map的计算复杂性,同时低分可以捕获更多的全局依赖性(构建分层特征图保持窗口大小不变)。对于局部来说,在transformer块的前馈网络的两个完全连接层引入深度卷积层,更好的捕获全局上下文
提出可学习的多尺度的恢复调解器(modulator)(如翻译不对请指正)*处理图像退化degardations。(以多尺度空间偏置形式调整Uformer解码器多层的特征,在每个LeWin Transformer块里会有一个可学习的基于窗口的Tensor张量被add到特征map里适应特征来获取更多的细节。)
方法Method
在这里插入图片描述
输入为一张degraded退化的图片(3HW),首先经过一个33的卷积(激活为leakyRelu)提取low-level特征X0,X0特征图经过K个编码器阶段,每个阶段stage包含一个LeWin Transformer块和一个下采样层。
(LeWinTransformer既有着可以获取长距离依赖的子注意力机制,也有着因为在特征图上没有重叠窗口而减少计算成本。)
(下采样层-将平铺的特征reshape成2D空间特征图,然后下采样map,用44,步长为2的卷积使得通道数2.也就是CH*
W->2CH/2W/2)
解码器同样经过k个stages阶段。和编码器一样,每个stage包括一个下采样层和一个LeWin Transformer。上采样是22转置卷积(非逆卷积,还原到原来大小,但数值不一样),步长为2。通道数减半,特征图size2,上采样后的特征和对应编码器特征进行cat后送入到LeWinTransformer进行图像恢复。K个解码器阶段后,将平铺的特征reshpe成2D特征图,经过33的卷积获得残差图像R(3H*W),最终增强的图像
I
′
=
I
+
R
\mathbf{I}^{\prime}=\mathbf{I}+\mathbf{R}
I′=I+R,训练Uformer使用Charbonnier 损失函数:
ℓ
(
I
′
,
I
^
)
=
∥
I
′
−
I
^
∥
2
+
ϵ
2
\ell\left(\mathbf{I}^{\prime}, \hat{\mathbf{I}}\right)=\sqrt{\left\|\mathbf{I}^{\prime}-\hat{\mathbf{I}}\right\|^{2}+\epsilon^{2}}
ℓ(I′,I^)=
I′−I^
2+ϵ2.一个是增强后的图像,一个是GT。
LeWin Transformer块
直接将transfomer运用到图像恢复具有两个挑战:1.在所有的tokens里计算全局自注意力会耗费巨大成本。尤其是在高分辨率上运动全局自注意力机制。2,全局上下文信息对于图像复原费油有必要,但是Transformer在捕获局部依赖时有限制。
基于这两个挑战,提出Locally-enhanced Window (LeWin) Transformer block,不仅可以因为Transformer的自注意力机制而获得长距离依赖,也可以因为将卷积放入Transfomer中获得有用的局部上下文信息。
X
l
′
=
W
−
M
S
A
(
LN
(
X
l
−
1
)
)
+
X
l
−
1
,
X
l
=
LeFF
(
LN
(
X
l
′
)
)
+
X
l
′
\begin{array}{l} \mathbf{X}_{l}^{\prime}=\mathrm{W}-\mathrm{MSA}\left(\operatorname{LN}\left(\mathbf{X}_{l-1}\right)\right)+\mathbf{X}_{l-1}, \\ \mathbf{X}_{l}=\operatorname{LeFF}\left(\operatorname{LN}\left(\mathbf{X}_{l}^{\prime}\right)\right)+\mathbf{X}_{l}^{\prime} \end{array}
Xl′=W−MSA(LN(Xl−1))+Xl−1,Xl=LeFF(LN(Xl′))+Xl′
W
−
M
S
A
\mathrm{W}-\mathrm{MSA}
W−MSA为非重叠的基于窗口的多头子关注
LeFF
\operatorname{LeFF}
LeFF为局部增强的前馈网络
W-MSA
用M*M的窗口大小分割输入2D特征X,从每个窗口i得到平铺的和转置的特征
X
i
∈
R
M
2
×
C
\mathbf{X}^{i} \in \mathbb{R}^{M^{2} \times C}
Xi∈RM2×C,然后我们在每个窗口的平铺特征上实行自注意力机制。
WQ,WK,WV分别代表q,k,v,
X
k
X_{k}
Xk是第k层head的输出,把k层head的输出cat到一起然后线性投射得到最后的结果。(当然,和之前的Transformer一样,也采用了位置编码。)
Locally-enhanced Feed-Forward Network (LeFF)
问题:标准Transformer中的前馈网络FFN在获取局部上下文有限制,但事实上在图像恢复的领域,相邻像素是很重要的参考信息。
解决:add一个深度卷积块在FFN里
在这里插入图片描述
3.3多尺度恢复调解器
为了校准特征map并鼓励恢复细节,提出多尺度恢复调解器。
多尺度恢复调解器将多个调解器运用在Uformer的解码器里,一个modulator调解器可以作为一个可学习的tensor-MxM*C(M是window size),每个modulator作为共享的bias偏置,添加到MSA自注意力模块前的所有非重叠窗口中。
(这个多尺度恢复调解器被作者验证在图像去噪和图像模糊的任务上很有效,可以去除更多的模糊噪声模式,产生更加清晰的图像)
实验
实验结果对比看的化,相比较其他论文工作,有提升,不明显。
作者相继做了去噪,运动去模糊,消除模糊,去雨的实验。
还有消融实验
表6 作者做了实验为了证明这个工作的分层结构。作者做了一个基于VIT架构仅包含单尺度的特征map。和Unet和Vit相比,我们的实验结果优于他们。
我看Uformer-B(Base)既有W-MSA和FFN效果还没只有FFN好–解释为什么将局部性引入前馈网络:相比于将局部性纳入W-MSA模块相比,将局部性引入前馈网络效果更好。
这篇贴出来了代码
https://github.com/ZhendongWang6/Uformer.