UNetFormer

为了充分利用图像信息,平衡局部与全局信息在语义分割中的提取。在Unet网络的基础上设计了以CNN提取图像局部特征的编码器和以Transformer为基础的增强图像细节特征和全局信息的解码器作为网络的主体结构。

网络结构

编码器

以四阶段ResBlock18为主干网络提取特征多尺度语义特征,四阶段过程中不断下采样,每个阶段都与解码器相连。

解码器

传统语义分割中捕获全局信息通常是在网络的最后阶段加上单个注意力模块或是使用Transformer构建编码器提取全局特征。前者不能充分捕获多尺度全局特征,后者提高了模型的计算复杂度。而在本网络中,使用了三个GLTB块和FRH来构建轻量级的解码器。

GLTB(Global-local Transformer Block)

GLTB将标准Transformer块中的多头注意力变为了全局-局部注意力。该部分采用双分支结构来提取融合全局和局部特征。

Global branch

全局分支采用基于窗口的多头注意力。

将输入采用1x1卷积增加三倍通道数,使用窗口划分,对于每个窗口reshape成一维后划分成Q,K,K输入到到分支中进行计算。

由于不同窗口的划分操作损失了场景的空间一致性,在最后设计了交叉窗口特征交互模块来恢复全局性。

交叉窗口采用水平池化和垂直池化来融合窗口特征,对于单个窗口而言,以$P_1$中红线为例,构建红线内的像素间的注意力关系,并构建水平方向上和垂直方向上其他窗口像素的关联。从而得到全局信息。

Local Branch

对于局部特征的提取,使用双分支的平行卷积操作。

FRH(Feature refinement head)

原来的Resblock中包含了丰富的空间信息,将这部分特征与全局-局部处理得到的信息经过权重融合输入。输入后使用双分支来分别增强通道信息和空间信息。通道分支中使用了全局平均池化,使用两层1x1卷积先降维再升维学习抽象特征。空间分支中使用深层卷积提取。最后将分支结果加法融合,加入残差连接防止网络衰减,最后上采样获得最终的语义图。