Swin Transformer + CNN

Abstract

卷积神经网络（CNN）推动了遥感图像的应用研究，由于感知场的大小固定，CNN无法建模全局语义相关性。利用基于自注意变压器的模型，可以对全局语义信息进行建模。然而，变压器用于自注意计算的斑块计算方法忽略了每个patch内部的空间信息。为了解决这些问题，我们提出了STransFuse模型作为一种新的遥感图像语义分割方法。该模型结合了变压器与CNN的优点，提高了各种遥感图像的分割质量，我们使用一个分阶段的模型来在不同的语义尺度上提取粗粒度和细粒度的特征表示，而不像早期基于变压器模型融合的技术。为了充分利用在不同阶段获得的特征，我们设计了一个自适应融合模块，该模块利用自注意机制，自适应地在不同尺度上的特征之间融合语义信息。我们提出的模型在 Vaihingen 数据集上的总体精度（OA）比基线高了1.36%，在 Potsdam 数据集上的OA比基线提高了1.27%.

本文主要贡献：

• A framework for Swin Transformer and CNN parallelism was designed with STransFuse model. 首先，STransFuse通过变换网络将特征图的全局语义信息引入到模型中，并利用CNN提取低层次特征图的空间上下文信息。其次，STransFuse 通过不需要创建非常深的网络，避免了梯度消失和特征图信息丢失的困难。然后，实验证明，所采用的双分支结构使模型在性能和计算速度方面都很有效。最后，设计的分阶段融合结构允许高级特征图包含更丰富的特征信息。
• 设计了一个特征映射融合模块，通过自适应融合特征映射与自注意结构来改进模型的特征表示。
• 所提出的 STransFuse 模型在 Vaihingen 和Potsdam 数据集上取得了较好的效果。

Method

A. Overview

本文提出了STransFuse模型作为一种新的语义分割方法，可以有效地获取遥感图像中的全局语义上下文信息和空间上下文信息。我们使用Swin变换器和CNN来处理图像，在不同的阶段融合特征图，最后将特征图恢复到它们的原始大小。

B. STransFuse Overall Architecture

为了提高模型分割的准确性，图像中的语义和空间上下文信息至关重要。CNN受到卷积核的固定大小的限制，不能建模全局语义信息。变换器可以通过自注意计算获得全局语义信息，但自注意计算需要将 patch 扩展成一维 token，patch 内部的空间信息是丢失的，为了解决这些问题，我们设计了具有两个编码器的模型来提取特征。我们使用具有训练权重的Resnet34作为CNN分支的编码器。通过融合从变压器分支和Resnet34分支中提取的特征，提高了模型的特征表示能力。

如图1(a)所示，将图像 $x\in R^{H\times W\times C}$ 分别输入Swin变压器网络和Resnet34网络，其中H表示图像的高度，W表示图像的宽度，C表示图像中的通道数。

模型在不同阶段提取的特征图大小不同，对应于特征粒度的不同语义尺度。为了使特征映射能够包含丰富的语义信息和特征细节信息，我们使用了一种分阶段的融合策略。

在Swin变压器网络中，有四个阶段分别获取 $x_{s_1},x_{s_2},x_{s_3},x_{s_4}$ 特征图，每个阶段都包含$patchmergingandSwinTransformer.$
补丁合并的工作方式类似于CNN的池化层，因为它会对图像进行降采样。补丁合并通过滑动输入图像上滑动窗口将图像分割成不重叠的补丁。每个补丁都被认为是一个“token”。我们最初将patch 的大小固定为$4\times 4$. 然后，通过线性嵌入层将特征图中的特征值投影到C维数上。最后，将Swin变压器块应用于这些补丁标记。输出特征图的分辨率为 $H/4\times W/4$. 上述步骤统称为“第一阶段”。在下面的“阶段2”中，补丁合并连接了每一组$2\times 2$ 个相邻补丁的特征，如图2所示。补丁合并采用线性嵌入层将输出尺寸改变为2C，并采用Swin变换器进行特征变换。在“阶段2”中，输出特征图的分辨率保持在 $H/8\times W/8$ 。“阶段3”和“阶段4”与“阶段2”相似，输出的特征图分辨率分别为$H/16\times W/16$ 和$H/32\times W/32$.

将图像输入Resnet34网络得到特征图，特征图由层4分别作为特征图$x_{c2},x_{c3}$输出到层4，这些特征图的大小分别为$H/8\times W/8,H/16\times W/16$ 和 $H/32\times W/32$。将Resnet34生成的特征图与Swin变压器不同阶段生成的特征图合并，利用Swin变压器收集特征的全局语义上下文信息的能力。最后，将融合后的特征图进行两次上采样，并将特征图恢复到输入图像的大小。

C. Swin Transformer Block

图像的上下文信息对于提高语义分割精度至关重要，而远程语义信息可以作为一种识别辅助手段，使模型不仅仅依赖于图像的光谱信息。因此，我们在模型中引入了变压器网络作为特征提取。有许多研究表明，变压器中的自我注意可以模拟全局语义信息[11]、[39]、[40]。

在标准变压器块中使用的自我注意是通过将其中一个令牌与所有其他令牌相关联来计算的。这种计算使网络的计算工作量相对于图像的分辨率大小呈二次增长，对于一些密集的预测任务（如语义分割），该模型将需要高端的计算设备。Swin变压器将在一个局部窗口中执行自注意计算。这些图像由不重叠的窗口进行分割。每个窗口都包含M×个M补丁。值得注意的是，为了确保图像（h×w）可以被窗口（M×M）整除，我们使用了一种填充方法，并在计算注意力时掩盖了填充值。在这种情况下，多头自注意（MSA）的计算复杂度如(1)所示，窗口MSA（W-MSA）的计算复杂度如(2)所示。

在(1) 中，MSA的计算复杂度与h和w的产生是二次的。在(2)中，当M是固定大小（默认设置为7）时，W-MSA的计算复杂度与h和w的产生呈线性相关。

Swin变压器是用W-MSA替换变压器模块中的MSA。如图1(b)所示，Swin变压器将拼接合并处理后的特征图输入到Swin变压器块中。然后，特征图通过LayerNorm层进入W-MSA模块，每个模块与另一个LayerNorm层之间存在残差连接。

由于W-MSA进行了滑动窗口分割操作，裁剪后的补丁不重叠，窗口之间缺乏有效的信息交互。存在一个移位窗口MSA（SW-MSA）网络，进一步提高了模型的性能。与W-MSA相比，SW-MSA执行窗口转换。SW-MSA的想法是将图像周期性地向上移动，并周期性地离开窗口大小的一半。在窗口之外的图像上的区域将分别移动到窗口的下方和右侧。然后，通过在移位之上根据W-MSA对窗口进行切片，我们将得到一种与W-MSA不同的窗口切片方法。SW-MSA的公式如下所示：

D. AFM Block

为了有效地融合了CNN和Swin变压器的编码特征，我们设计了一个基于自注意机制的AFM。通过自注意计算得到特征权矩阵，选择性地增强空间细节或抑制其他区域，从而提高了密集预测的分化能力。AFM的结构如图3所示。我们将用以下的方程式来进行特征的融合：

其中，$x_{f,i}$ 表示AFM第 $i$ 级输出的特征矩阵，$x_{c,i}$ 表示CNN第$i$ 层输出的特征矩阵，$x_{s,i}$ 表示Swin变压器第$i$ 级输出的特征矩阵，$x_{BN,i}$ 表示第$i$ 个AFM块融合特征图，$x_q$ 是自注意计算中的query.

由于使用了三个分支的特征映射，特征融合计算密集型。为了缓解这一问题，我们在AFM中添加了自适应aviavgpool2d方法，使AFM能够构建每个像素和一些收敛中心之间的关系。通过从输入张量中像素子集的特征向量集合，AFM在计算上是可接受的。AFM中语义相关性的计算公式如下所示：

$x_{f,i}$中每个位置的特征值是所有位置的特征与原始特征的加权之和。因此，AFM能够基于全局视图中的注意特征图，选择性地聚合上下文信息，提高模型区分密集的像素的能力。每个像素的信息将传递给与其语义相关的像素，提高语义一致性。