【遥感语义分割】UNetFormer

原文：UNetFormer: An UNet-like Transformer for Efficient Semantic Segmentation of Remotely Sensed Urban Scene Imagery

Libo Wang1, 2, Rui Li1, Ce Zhang3, 4, Shenghui Fang1*, Chenxi Duan5, Xiaoliang Meng1, 2 and Peter M. Atkinson3, 6, 7

1) 中国湖北省武汉市珞瑜路129号，武汉大学遥感信息工程学院。
2) 中国广东省广州市自然资源部南亚热带和亚热带地区自然资源监测重点实验室。
3) 英国兰开斯特大学兰开斯特环境中心，邮政编码LA1 4YQ。
4) 英国生态与水文学中心，兰开斯特，邮政编码LA1 4AP。
5) 荷兰特文特大学地理信息科学与地球观测学院。
6) 英国南安普顿大学地理与环境科学学院，南安普顿Highfield校区，邮政编码SO17 1BJ。
7) 中国科学院地理科学与资源研究所，北京市大屯路甲11号，邮政编码100101。
*通讯作者。

摘要——遥感城市场景图像的语义分割在土地覆盖图绘制、城市变化检测、环境保护和经济评估等广泛的实际应用中都有需求。在深度学习技术的快速发展推动下，卷积神经网络（CNN）在多年来一直主导着语义分割。CNN采用分层特征表示，展示了强大的局部信息提取能力。然而，卷积层的局部特性限制了网络捕捉全局上下文信息的能力。最近，作为计算机视觉领域的热点，Transformer展示了其在全局信息建模方面的巨大潜力，提升了许多与视觉相关的任务，如图像分类、目标检测，特别是语义分割。在本文中，我们提出了一种类似UNet的Transformer（UNetFormer）用于实时的城市场景分割。新颖的UNetFormer采用了一种基于CNN的编码器和基于Transformer的解码器的混合结构，以高计算效率学习全局-局部上下文。大量实验表明，所提出的UNetFormer不仅在推理阶段运行更快，而且相比于最新的轻量级模型还具有更高的精度。具体来说，所提出的UNetFormer在UAVid测试集上达到了67.8%的mIoU，在LoveDA数据集上达到了52.4%的mIoU，而在NVIDIA GTX 3090 GPU上，输入为512×512时，推理速度可达到每秒322.4帧。源码将在https://github.com/WangLibo1995/UNetFormer免费提供。

关键词——语义分割, 遥感, 视觉Transformer, 混合结构, 全局-局部上下文, 城市场景。

1. 引言
随着传感器技术的进步，全球范围内越来越多地捕获到高分辨率的遥感城市场景图像，这些图像具有丰富的空间细节和潜在的语义内容。城市场景图像广泛应用于语义分割任务，即像素级别的分割和分类，进而应用于各种与城市相关的领域，包括土地覆盖图绘制（Li等, 2022b；Maggiori等, 2016；Marcos等, 2018）、变化检测（Xing等, 2018；Yin等, 2018）、环境保护（Samie等, 2020）、道路和建筑提取（Griffiths和Boehm, 2019；Shamsolmoali等, 2020；Vakalopoulou等, 2015）以及许多其他实际应用（Picoli等, 2018；Shen等, 2019）。最近，深度学习技术（LeCun等, 2015），尤其是卷积神经网络（CNN）的兴起，主导了语义分割任务（Chen等, 2014；Chen等, 2018b；Long等, 2015；Ronneberger等, 2015；Zhao等, 2017a）。与传统的分割机器学习方法相比，如支持向量机（SVM）（Guo等, 2018）、随机森林（Pal, 2005）和条件随机场（CRF）（Krähenbühl和Koltun, 2011），基于CNN的方法能够捕捉到更精细的局部上下文信息，这支持了其在特征表示和模式识别方面的强大能力（Zhang等, 2020a；Zhang等, 2020b）。

尽管上述具有优势，但在对高分辨率遥感城市场景图像进行标注时，基于CNN的方法面临的主要挑战是全局上下文信息的提取。具体来说，CNN设计用于提取局部模式，天然缺乏建模全局上下文的能力。然而，如图1所示，全局和局部信息的交互对于区分城市中复杂且异质的人造物体都是至关重要的（Yang等, 2021a）。尽管将最新的注意力机制融入CNN在一定程度上缓解了上述问题，但通常需要大量计算时间和内存来捕捉全局上下文（Li等, 2021c），从而降低了它们的效率，并限制了它们在实时城市应用中的潜力。

图1. 全局和局部上下文的示意图。

本文旨在在确保网络效率的同时，实现精确的城市场景分割。受Transformer在计算机视觉领域最新突破的启发，我们提出了一种类似UNet的Transformer（UNetFormer）来应对这一挑战。UNetFormer创新性地采用了基于CNN的编码器和基于Transformer的解码器的混合架构。具体来说，我们采用ResNet18作为编码器，并设计了一个全局-局部Transformer块（GLTB）来构建解码器。与标准Transformer中的传统自注意力块不同，所提出的GLTB开发了一种高效的全局-局部注意力机制，包含一个注意力全局分支和一个卷积局部分支，以捕捉视觉感知的全局和局部上下文，如图2所示。在全局分支中，引入了基于窗口的多头自注意力概念，以低复杂度捕捉全局上下文（Liu等, 2021）。随后，为了加强全局和局部上下文之间的交互，我们开发了一个十字形上下文交互模块。最后，为了有效融合空间细节和上下文信息以及进一步优化特征图，提出了一个特征优化头（FRH），并将其附加在网络末端。精度和效率之间的权衡以及有效的特征优化使得所提出的方法在四个公共数据集上，UAVid（Lyu等, 2020）、ISPRS Vaihingen和Potsdam数据集以及LoveDA（Wang等, 2021a），超越了最新的轻量级网络，实现了高效的遥感城市场景图像分割。

本文的其余部分组织如下：在第2部分中，我们回顾了基于CNN和基于Transformer的城市场景分割以及全局上下文建模的相关工作。在第3部分中，我们展示了UNetFormer的结构，并介绍了所提出的GLTB和FRH。在第4部分中，我们进行消融研究，以证明GLTB和FRH以及新颖的混合结构的有效性，并将结果与应用于四个数据集的一组最新模型进行比较。在第5部分中，我们提供了全面的讨论。第6部分是总结和结论。

图2 (a) 标准Transformer块示意图和 (b) 全局-局部Transformer块示意图。

2. 相关工作
2.1 基于CNN的语义分割方法
全卷积网络（FCN）（Long等, 2015）是第一个以端到端方式解决语义分割问题的有效CNN结构。自那以后，基于CNN的方法在遥感领域的语义分割任务中占据了主导地位（Kemker等, 2018；Kotaridis和Lazaridou, 2021；Ma等, 2019；Tong等, 2020；Zhao和Du, 2016；Zhu等, 2017）。然而，FCN的过于简化的解码器导致分割结果分辨率较低，限制了其精细度和准确性。

为解决这一问题，提出了一种编码器-解码器网络，即UNet，用于语义分割，包括两条对称路径，称为收缩路径和扩展路径（Ronneberger等, 2015）。收缩路径通过逐渐降低特征图的空间分辨率来提取分层特征，而扩展路径则通过逐步恢复空间分辨率来学习更多的上下文信息。随后，编码器-解码器框架已成为遥感图像分割网络的标准结构（Badrinarayanan等, 2017；Chen等, 2018a）（Sun等, 2019）。基于编码器-解码器结构，（Diakogiannis等, 2020；Yue等, 2019；Zhou等, 2018）设计了不同的跳跃连接以捕捉更丰富的上下文信息，而（Liu等, 2018；Zhao等, 2017b）（Shen等, 2019）开发了各种解码器以保留语义信息。

基于编码器-解码器的CNN方法，尽管取得了令人鼓舞的性能，但在城市场景解释中遇到了瓶颈（Sherrah, 2016）（Marmanis等, 2018；Nogueira等, 2019）。具体来说，具有有限感受野的CNN分割网络只能提取局部语义特征，缺乏从整个图像中建模全局信息的能力。然而，在高分辨率的遥感城市场景图像中，复杂的模式和人造物体经常出现（Kampffmeyer等, 2016；Marcos等, 2018）（Audebert等, 2018）。如果仅依赖局部信息，很难识别这些复杂的对象。

2.2 全局上下文建模
为了使网络摆脱CNN局部模式的焦点，许多尝试都集中在建模全局上下文信息上，其中最流行的方法是将注意力机制整合到网络中。例如，Wang等修改了点积自注意力机制，并将其应用于计算机视觉领域（Wang等, 2018）。Fu等在扩展FCN顶部附加了两种类型的注意力模块，以自适应地整合局部特征与其全局依赖关系（Fu等, 2019）。Huang等提出了一个十字形注意力块来聚合信息丰富的全局特征（Huang等, 2020）。Yuan等开发了一个对象上下文块，探索基于对象的全局关系（Yuan等, 2020）。

注意力机制也提升了遥感图像分割网络的性能。例如，Yang等提出了一种融合高级和低级语义特征的注意力融合网络，并在高分辨率遥感图像语义分割中取得了最先进的结果（Yang等, 2021b）。Li等集成了轻量级的空间和通道注意力模块，以自适应方式优化高分辨率遥感图像的语义特征（Li等, 2020a）。Ding等设计了一个带有嵌入模块的局部注意力块，以捕捉更丰富的上下文信息（Ding等, 2021）。Li等开发了一种线性注意力机制，以降低计算复杂度同时提高性能（Li等, 2021a）。然而，上述注意力模块由于过度依赖卷积操作而限制了全局特征表示。此外，单一的注意力模块无法在解码器中多层语义特征上建模全局信息。

2.3 基于Transformer的语义分割方法
近年来，有几种尝试将Transformer应用于全局信息提取（Vaswani等, 2017）。不同于CNN结构，Transformer将基于2D图像的任务转换为基于1D序列的任务。由于强大的序列到序列建模能力，Transformer展示出比仅使用注意力的模型更优越的全局上下文提取能力，并在基础视觉任务如图像分类（Dosovitskiy等, 2020）、目标检测（Zhu等, 2020）和语义分割（Zheng等, 2021）等方面取得了最先进的结果。在此推动下，许多遥感领域的研究人员已将Transformer应用于遥感图像场景分类（Bazi等, 2021；Deng等, 2021）、高光谱图像分类（Hong等, 2021；He等, 2021；Zhong等, 2021）、目标检测（Li等, 2022a）、变化检测（Chen等, 2021a）、建筑和道路提取（Chen等, 2021c）（Sun等, 2022），特别是语义分割（Li等, 2020b）（Gao等, 2021）（Zhang等, 2022）。

目前大多数用于语义分割的Transformer仍然遵循编码器-解码器框架。根据不同的编码器-解码器组合，它们可以分为两类。第一类由基于Transformer的编码器和基于Transformer的解码器构成，即纯Transformer结构。典型模型包括Segmenter（Strudel等, 2021）、SegFormer（Xie等, 2021）和SwinUNet（Cao等, 2021）。第二类采用混合结构，由基于Transformer的编码器和基于CNN的解码器组成。基于Transformer的语义分割方法通常采用第二种结构。例如，TransUNet采用了混合视觉Transformer（Dosovitskiy等, 2020）作为编码器进行更强的特征提取，并在医学图像分割中取得了最先进的结果（Chen等, 2021b）。DC-Swin引入了Swin Transformer（Liu等, 2021）作为编码器，并设计了一个密集连接的卷积解码器用于高分辨率遥感图像分割，大幅超越了基于CNN的方法（Wang等, 2022）。Panboonyuen等（2021）也选择了Swin Transformer作为编码器，并利用了UNet（Ronneberger等, 2015）、FPN（Kirillov等, 2019）和PSP（Zhao等, 2017a）等各种基于CNN的解码器，用于遥感图像的语义分割，取得了先进的准确性（Panboonyuen等, 2021）。

尽管上述优点，由于其平方复杂度的自注意力机制（Vaswani等, 2017），基于Transformer的编码器的计算复杂度远高于基于CNN的编码器，这严重影响了它在城市相关实时应用中的潜力和可行性。因此，在本文中，为了充分利用Transformer的全局上下文提取能力而不带来高计算复杂度，我们提出了一种UNet-like Transformer结构，其中采用了基于CNN的编码器和基于Transformer的解码器，用于高效地对遥感城市场景图像进行语义分割。具体而言，对于我们的UNetFormer，我们选择了轻量级骨干网络ResNet18作为编码器，并开发了高效的全局-局部注意力机制来构建解码器中的Transformer块。所提出的全局-局部注意力机制采用了双分支结构，即全局分支和局部分支。这种结构使注意力块能够捕捉全局和局部上下文，从而超越了仅捕捉全局上下文的单分支高效注意力机制（Liu等, 2021；Zhang和Yang, 2021）。

3. 方法论
如图3所示，所提出的UNetFormer采用了基于CNN的编码器和基于Transformer的解码器进行构建。接下来的各节将详细描述每个组件。

Fig. 3. UNetFormer概述。

3.1 基于CNN的编码器
作为ResNet18（He等, 2016）在广泛的实时语义分割任务中同时展示了有效性和效率，我们选择预训练的ResNet18作为编码器，以显著降低计算成本提取多尺度语义特征。ResNet18由四个阶段的Resblock组成，每个阶段通过2的尺度因子对特征图进行下采样。在提出的UNetFormer中，每个阶段生成的特征图通过通道维度为64的1×1卷积与解码器的相应特征图融合，即跳跃连接。具体来说，Resblocks生成的语义特征与解码器的GLTB生成的特征使用加权和操作聚合。加权和操作根据它们对分割准确性的贡献有选择地加权两个特征，从而学习更加泛化的融合特征（Tan等, 2020）。加权和操作的表述可以表示为：
$\mathbf{FF} = \alpha \cdot \mathbf{RF} + (1 - \alpha) \cdot \mathbf{GLF}$
其中，FF表示融合后的特征，RF表示Resblocks生成的特征，GLF表示全局-局部Transformer块生成的特征。

3.2 基于Transformer的解码器
在细分辨率的遥感城市图像中，复杂的人工制品经常出现，这使得在缺乏全局语义信息的情况下难以实现精确的实时分割。为了捕捉全局上下文，主流解决方案通常集中在网络末端附加单个注意力块（Wang等, 2018）或引入Transformer作为编码器（Chen等, 2021b）。前者不能捕捉多尺度的全局特征，而后者显著增加了网络的复杂性并丢失了空间细节。相比之下，在提出的UNetFormer中，我们利用三个全局-局部Transformer块和一个特征细化头部构建了一个轻量级的基于Transformer的解码器，如图3所示。通过这种分层和轻量化的设计，解码器能够在多个尺度上捕捉全局和局部上下文，同时保持高效率。

3.2.1 全局-局部Transformer块 (GLTB)
全局-局部Transformer块由全局-局部注意力、多层感知器、两个批量归一化层和两个附加操作组成，如图2(b)所示。
全局-局部注意力：虽然全局上下文对于复杂城市场景的语义分割至关重要，但局部信息对于保留空间细节仍然必不可少。在这方面，所提出的全局-局部注意力构建了两个并行分支，分别提取全局和局部上下文，如图4(a)所示。

图4. 全局-局部注意力的示意图。

作为一个相对浅层的结构，局部分支使用了两个并行的卷积层，卷积核尺寸分别为3和1，用于提取局部上下文。然后在最终求和操作之前附加了两个批量归一化操作。

全局分支部署了基于窗口的多头自注意力机制来捕捉全局上下文。如图4(b)所示，我们首先使用标准的1×1卷积将输入二维特征图 $\in \mathbb{R}^{B \times C \times H \times W}$ 的通道维度扩展到三倍。然后，我们应用窗口分割操作，将一维序列 $\mathbb{R}^{\left(3 \times B \times \frac{H}{w} \times \frac{W}{w} \times h\right) \times(w \times w) \times \frac{C}{h}}$ 切分成查询（Q）、键（K）和值（V）向量。通道维度 $C$ 被设置为64，窗口大小w和头的数量 h都设置为8。基于窗口的多头自注意力机制的详细信息可以参考Swin Transformer (Liu et al., 2021)。

在非重叠的局部窗口中执行自注意力机制，虽然高效，但由于缺乏跨窗口的交互，可能会破坏城市场景的空间一致性。Swin Transformer 引入了一个额外的移位 Transformer 块来挖掘局部窗口之间的关系。尽管捕捉跨窗口关系的能力有所提高，但计算量也显著增加。在本文中，我们提出了一种交叉形上下文交互模块，以高计算效率捕捉跨窗口关系。如图5所示，交叉形上下文交互模块融合了由水平平均池化层和垂直平均池化层生成的两个特征图，以加强窗口内的关系。此外，为了利用局部语义信息的优势，融合后的交叉形上下文进一步与局部上下文聚合。最后，我们采用深度卷积、批归一化操作和标准的1×1卷积来表征细粒度的全局-局部上下文。

图 5. 交叉形上下文交互模块。

3.2.2 特征优化头 (FRH)
第一个 Resblock 生成的浅层特征保留了城市场景丰富的空间细节，但缺乏语义内容；而深层全局-局部特征提供了精确的语义信息，但空间分辨率较粗糙。因此，尽管直接对这两个特征进行求和操作速度快，但可能会降低分割精度（Poudel 等人，2018；Poudel 等人，2019；Yu 等人，2018）。在本文中，我们开发了一个特征优化头，用于缩小这两个特征之间的语义差距，以进一步提高精度。

如图 6 所示，首先对这两个特征进行加权求和操作，以充分利用精确的语义信息和空间细节。融合特征随后被选为 FRH 的输入，如图 3 所示。其次，我们构建两条路径来加强通道和空间特征表示。具体来说，通道路径利用全局平均池化层生成通道注意力图 $\mathbf{C} \in \mathbb{R}^{1 \times 1 \times c}$ ，其中 $c$ 表示通道维度。减少和扩展操作包含两个 1×1 卷积层，首先将通道维度 $c$ 减少四分之一，然后再将其扩展回原始尺寸。空间路径利用深度卷积生成空间注意力图 $\mathbf{S} \in \mathbb{R}^{h \times w \times 1}$ ，其中 $h$ 和 $w$ 表示特征图的空间分辨率。两条路径生成的注意力特征通过求和操作进一步融合。最后，应用后处理的 1×1 卷积层和上采样操作生成最终的分割地图。值得注意的是，引入了残差连接以防止网络退化。

图6. 特征细化头（Feature Refinement Head）。

3.3 损失函数

在训练阶段，我们不仅使用主要的特征细化头，还构建了一个额外的辅助头来优化全局局部Transformer块，如图3所示。这种多头分割架构在先前的研究中已被证明是有效的（Yu et al., 2020; Zhu et al., 2019）。基于多头设计，我们应用一个主要损失和一个辅助损失来训练整个网络。主要损失 $\mathcal{L}_{p}$ 是Dice损失 $\mathcal{L}_{\text{dice}}$ 和交叉熵损失 $\mathcal{L}_{ce}$ 的组合，可以表示为：

$\begin{array}{c} \mathcal{L}_{ce} = -\frac{1}{N} \sum_{n=1}^{N} \sum_{k=1}^{K} y_{k}^{(n)} \log \hat{y}_{k}^{(n)} \\ \mathcal{L}_{\text{dice}} = -\frac{2}{N} \sum_{n=1}^{N} \sum_{k=1}^{K} \frac{\hat{y}_{k}^{(n)} y_{k}^{(n)}}{\hat{y}_{k}^{(n)} + y_{k}^{(n)}} \\ \mathcal{L}_{p} = \mathcal{L}_{ce} + \mathcal{L}_{\text{dice}} \end{array}$

其中， $N$ 和 $K$ 分别表示样本数量和类别数量。 $y^{(n)}$ 和 $\hat{y}^{(n)}$ 分别表示真实语义标签的独热编码和网络输出的softmax结果， $n \in [1, \cdots, N]$ 。 $\hat{y}_{k}^{(n)}$ 表示样本 $n$ 属于类别 $k$ 的置信度。我们选择交叉熵损失作为辅助损失 $\mathcal{L}_{aux}$ ，并将其应用于辅助头部。辅助头部以三个全局局部Transformer块的融合特征作为输入，构建了一个包含批标准化和ReLU的 $3 \times 3$ \卷积层，一个 $1 \times 1$ 卷积层以及一个上采样操作来生成输出。为了更好地结合主要损失，辅助损失还会乘以一个因子 $\alpha$ 。因此，整体损失 $\mathcal{L}$ 可以表示为：