UNetFormer

UNetFormer: A UNet-like transformer for efficient semantic segmentation of remote sensing urban scene imagery
UNetFormer：一种用于高效遥感城市场景图像语义分割的类UNet变压器模型
论文地址
代码地址

1. 介绍

• 轻量级设计：UNetFormer使用了ResNet18作为编码器。
• 全局-注意力机制：通过三个全局-局部Transformer块在解码器中同时捕获多尺度的全局和局部上下文信息。
• 特征精细化头部（FRH）：通过加权融合浅层的空间细节和深层语义信息，进一步提升分割精度。

这篇论文的基础网络结构是U-Net，编码器使用ResNet18，解码器修改为设计的模块，主要还是使用Transformer，切入点是如果直接使用Transformer的话，会导致计算复杂度增加并且会丢失空间细节信息，因此作者设计的模块通过全局-局部注意力机制模块，解决了计算复杂度、全局上下文的信息和空间细节方面的不足。

2. 网络结构

2.1 CNN-based encoder

编码器阶段采用预训练的ResNet18，编码后的特征图通过1 $\times$ 1的卷积与解码器对于阶段的特征图进行融合，也就是跳跃连接。具体来说，ResBlock生成的语义信息通过加权求和与解码器阶段的GLTB模块生成的特征进行融合，加权求和操作的公式可以表示为： $FF=\alpha \cdot RF+(1-\alpha )\cdot GLE$，其中，$FF$表示融合后的特征，$RF$表示由Resblocks生成的特征，$GLF$表示由全局-局部Transformer块生成的特征。

2.2 Transformer-based decoder

在UNetFormer中，利用了三个全局-局部Transformer块和一个特征精细化头部构建了一个轻量级的基于Transformer的解码器，接下来细看解码器的各个部分。

2.2.1 Global-local Transformer Block (GLTB）

上图是全局-局部Transformer块，由全局-局部注意力、多层感知机、两个批量归一化层以及两个附加操作组成。主要的变化是在Efficient Global-local Attention模块，这个模块包含了提取全局信息的多头窗口自注意力机制和提取局部信息的普通卷积。

上图是作者提出的全局-局部注意力构建了两个分支，分别提取全局和局部上下文信息。
局部分支使用两个并行卷积层（核大小分别为3和1）来提取局部上下文信息，随后，在最终的求和操作之前附加了两个批量归一化操作。全局分支则采用基于窗口的多头自注意力机制来捕获全局上下文信息，如下图所示：

这里使用标准的1 $\times$ 1卷积将输入的特征图通道维度变为原来的3倍，接着通过窗口划分操作（window partition），单独计算各个窗口的注意力得分，这样就减少了计算的复杂度，对于每个单独的窗口，将去展平为一维序列单独计算它的Q、K、V向量（这里将通道维度设置为64，窗口大小和头数均设置为8）。对于各个窗口之间的关系，本文提出一种十字形窗口上下文交互模块，以高计算效率捕获跨窗口关系，如下图所示：

十字形窗口上下文交互模块在之前的一篇论文也看到过，这篇论文提出来的并且解释得特别详细。假设对于一个二维矩阵来说，十字形窗口上下文交互就是要捕获以某个点为中心，它的所在行所在列的特征信息，该模块融合了水平平均池化层和垂直平均池化层生成的两个特征图从而捕获全局上下文信息。具体而言，下面以水平平均池化层建立了窗口之间的水平关系为例，例如Win1=H(Win2)，对于窗口1中的任意点$P_1^{(m,n)}$，其与窗口2中点 $P_2^{(m+w,n)}$的依赖关系可以建模为：

其中w为窗口大小，D表示自注意力计算，可以建模局部窗口中像素对的依赖关系，因此，对于窗口1红色路径中的任意其他点$P_1^{(m+i,n)}$，其与$P_1^{(m,n)}$的依赖关系可以通过公式(3)建模；对于窗口2绿色路径中的其他任意其他点$P_2^{(m+w-j,n)}$，其与$P_2^{(m+w,n)}$的依赖关系可以通过公式(4)建模。公式(2)可以重写为公式(5)，即$P_1^{(m,n)}$与$P_2^{(m+w,n)}$的依赖关系被建模。基于这种跨窗口的像素级依赖关系，窗口1和窗口2之间的水平关系得以建立。类似地，窗口1和窗口3之间的垂直关系也可以通过相同方式建立，即$Win1=V(Win3)$，对于窗口4，则有$Win1=V(H(Win4)+H(V(Win4)))$。（这里的V表示水平方向，H表示垂直方向）
推广到$M\times M$输入（M表示窗口数量），通过连接更多中间窗口（如窗口2和窗口3），可以建模任意两个窗口之间的长程依赖关系。一次，十字形窗口上下文交互建模块能够建模窗口见的长程依赖关系，从而捕获全局上下文信息。
此外，全局分支中的全局上下文信息进一步与局部分支中的局部上下文信息聚合，生成全局-局部上下文，最后，采用深度可分离卷积、批量归一化操作以及标准$1\times 1$卷积来表征细粒度的全局-局部上下文信息。

2.2.2 Feature refinement head (FRH)

特征精细化头部：

我们在GLTB模块中，将全局信息和局部信息进行融合后，作为输入传入FRH模块，在这个模块，构建了两条路径来增强通道维度和空间维度的特征的表示能力。

• 对于通道路径：这里使用全局平均池化生成了一个注意力图$C\in R^{1\times 1\times c}$。所见与扩展操作包含两个$1\times 1$卷积层，首先将通道维度缩减4倍，在扩展恢复到原始维度。
• 对于空间路径：利用深度可分卷积生成空间注意力图$S\in R^{h\times w\times 1}$。

两条路径生成的注意力特征通过求和操作进一步融合。最后，应用一个后处理的$1\times 1$卷积层和上采样操作，生成最终的分割图。

3. 实验复现：

首先来看分割的结果图：

正好之前使用SFANet进行过分割，这里可以做一个对比：（左边是SFANet的分割效果，右边是UNetFormer的分割效果）

可以看出UNetFormer的分割效果更好，接下来看一个各项指标：

在训练集上的平均交并比是0.839，在验证集上的平均交并比是0.822，和SFANet做一下对比：

在SFANet上，训练集和验证集上的平均交并比分别为：0.879和0.83，比UNetFormer率高。