代码:https://github.com/ZhendongWang6/Uformer.

1.摘要

        本文提出一种运用于图像恢复的Transformer,本文取名为Uformer。其中我们使用Transformer块构建了一种类似于U-net的分层编码器-解码器网络。在Uformer中,有两个核心设计。首先,我们引入了一种新的局部增强窗口(LeWin) Transformer块,它执行基于非重叠窗口的自关注,而不是全局自关注。在获取局部上下文的同时,显著降低了高分辨率特征图的计算复杂度。其次,我们提出了一种多尺度空间偏置形式的可学习多尺度恢复调制器,用于调整Uformer解码器的多层特征。在这两种设计的支持下,Uformer具有捕获本地和全局依赖关系以进行图像恢复的高能力。为了评估我们的方法,在几个图像恢复任务上进行了大量的实验,包括图像去噪、运动去模糊、离焦去模糊和去训练。

2.本文方法

        本文的创新主要集中在两个方面,第一点是提出了一种改善的Transformer模块,然后运用到一种类似于Unet的结构中,第二点三在框架中加入了多尺度恢复调制器。下面将从总结构,Transformer模块和多尺度恢复调制器这三个方面详细介绍本文。完成本文的理解需要详细了解。

2,1总体结构

        如下图是我们网络的主要结构:

        如图所示,所提出的Uformer的整体结构是一个类似Unet的分层网络。具体来说,给定图像

I\in R^{3\times H\times W},

Uformer首先使用一个带有LeakyReLU的3× 3卷积层来提取底层特征。

X_{0}\in R^{C\times H\times W},

接下来,按照Unet的设计,将特征X_{0}通过K个编码器。每个阶段包含一堆LeWinTransformer块和一个下采样层。LeWinTransformer块利用了自关注机制来捕获远程依赖关系。

3.2. LeWinTransformer

        现在在应用Transformer进行图像恢复时有两个主要的挑战。

        (1)标准Transformer架构在所有像素之间计算全局关注,这导致了相对于像素数量的二次计算成本。在高分辨率特征图上应用全局自关注会导致计算成本过高,这显然是难以在实际情况上运用的。

       (2)局部上下文信息对于图像恢复任务至关重要,因为可以利用退化像素的邻域来恢复其原图像,但先前的工作表明,Transformer在捕获局部依赖项方面存在限制。

        为了解决上面提到的两个问题,我们提出了一个局部增强窗口(LeWin) Transformer块,如图2(b)所示,它受益于Transformer中的自关注来捕获远程依赖关系,并且还将卷积运算符包含到Transformer中以捕获有用的局部上下文。

LeWin Transformer块的计算表示为:

X_{l}'=W-MSA(LN(X_{l-1}))+X_{l-1},

X'=LeFF(LN(X_{l}'))+X_{l}'

LN表示layer normalization。下面我们分别对W-MSA和LeFF进行阐述。

Window-based Multi-head Self-Attention (W-MSA)

        本文没有像vanilla Transformer那样使用全局自关注,而是在不重叠的局部窗口内执行自关注,这大大降低了计算成本。给定2D特征地图X\in R^{C\times H\times W}, H和W分别为maps的高度和宽度,我们将X分割为不重叠的窗口,窗口大小为M×M,然后从每个窗口中获得flattend和转置的特征X_{i}\in R^{M^{2}\times C}。接下来,我们对每个窗口中的平坦化特征进行自关注。

        假设Transformer中的Heads数为k,每一个head的维数为d_{k}=C/k。则计算非重叠窗口内的第k个head的自关注可表示为:

X=\left \{ X^{1},X^{2},\cdots ,X^{N} \right \},N=HW/M^{2},

Y_{k}^{i}=Attention(X^{i}W_{k}^{Q},X^{i}W_{k}^{K},X^{i}W_{k}^{V}),i=1,\cdots ,N

\tilde{X}_{k}=\left \{ Y_{k}^{1}, Y_{k}^{2},\cdots Y_{k}^{N}\right \},

        其中,W_{k}^{Q},W_{k}^{Q},W_{k}^{Q} \in R^{C \times d_{k}}分别表示第k个heads的查询、键和值的矩阵。\tilde{X}_{k}是第k个头的输出。然后将所有头{1,2,···,k}的输出连接起来,然后进行线性投影以获得最终结果。我们也将相对位置编码应用到注意模块中(这里没有完全搞清楚,需要阅读参考文献43,51),注意计算公式为:

Attention(Q,K,V)=SoftMax(\frac{QK^{T}}{\sqrt{d_{k}}}+B)V,

        式中,B为相对位置偏差,其值取自\hat{B}\in R^{(2M-1) \times (2M-1)},其中的参数可学习[43,51]。与全局自注意相比,基于窗口的自注意可以显著降低计算成本。

        给定特征映射X\in R^{C\times H\times W},计算复杂度从O(H^{2}W^{2}C)下降到O(\frac{HW}{M^{2}}M^{4}C)=O(M^{2}HWC)。由于我们将Uformer设计为一个分层架构,我们基于窗口的自关注在低分辨率特征图上工作于更大的接受域,并且足以学习长期依赖关系。我们还在框架中每个阶段的even LeWin Transformer块中尝试了移动窗口策略[43],结果仅略好一些。

Locally-enhanced Feed-Forward Network (LeFF).

        本文使用的MLP模块与之前的区别就是使用了深度卷积,并没有特别难理解的地方,总体如图:

3.3 Multi-Scale Restoration Modulator (多尺度恢复调制器)

        不同类型的图像退化(如模糊、噪音、雨水等)有其独特的扰动模式需要处理或恢复。为了进一步提高Uformer处理各种扰动的能力,我们提出了一种轻量级的多尺度恢复调制器来校准特征并鼓励恢复更多的细节。如图2(a)和2(c)所示,多尺度恢复调制器在Uformer解码器中使用多个modulator。特别地,在每个LeWin Transformer块中,将调制器表示为形状为M×M×C的可学习张量,其中M为窗口大小,C为当前特征映射的通道维数。每个调制器简单地作为一个共享偏置项,在自注意模块之前添加到所有非重叠窗口中。更详细需要阅读[26]

这几天阅读了这篇文章,收获颇丰。但其中还有理解不足的地方:

(1)对位置偏差中的B矩阵具体是怎么取的,然后具体是怎么操作的。

(2)计算复杂度是怎么计算的。

(3)移动窗口策略是怎么实现的。

这三点就是本文阅读完之后的一些没理解的地方,我在此记录,希望通过之后的学习完善这些问题。

本文的启发:

通过对本文的阅读,为了解到在图像处理的问题中,transformer的改造主要在两个方面:

(1)减少计算的复杂度。

(2)提升tansfomer模块注意力机制对于各个像素的运用。

        为了解决这两个问题,所以transformer的改进集中在对transfomer模型的输入还有attention值的计算。还有tansformer模块的中MLP网络的修改。

        除了以上之外,还可以通过一些其他网络中迁移的小技巧来提升网络的性能,就比如本文中的Modular。

        需要详细的了解本文的计算方法在读完论文后还需要对代码进行详细了解。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐