CM-UNet: 利用Mamba架构的高效全局上下文建模进行图像语义分割

由于图像规模大和目标变化多样，当前的基于卷积神经网络（CNN）和基于Transformer的遥感图像语义分割方法在捕捉长距离依赖性方面存在不足，或者受限于复杂的计算复杂性。 在本文中，作者提出了CM-UNet，它包括一个基于CNN的编码器用于提取局部图像特征，以及一个基于Mamba的解码器用于聚集和整合全局信息，从而促进遥感图像的有效语义分割。具体来说，引入了CSMamba块来构建核心分割解码器，该块采用通道和空间注意力作为普通Mamba的门控激活条件，以增强特征交互和全局-局部信息融合。此外，为了进一步细化CNN编码器输出的特征，采用了一个多尺度注意力聚合（MSAA）模块来合并不同尺度的特征。 通过整合CSMamba块和MSAA模块，CM-UNet有效地捕捉了大尺度遥感图像的长距离依赖性和多尺度全局上下文信息。在三个基准数据集上获得的实验结果表明，所提出的CM-UNet在各种性能指标上均优于现有方法。 代码可在https://github.com/XiaoBuL/CM-UNet获取。

I Introduction

远程遥感图像语义分割涉及将大规模遥感图像中的像素分类到不同的类别中，以增强对遥感（RS）数据的分析和解释。这种大规模的语义分割对于自动驾驶[1]、城市规划[2]、环境保护[3]以及其他许多实际应用都至关重要。

随着深度学习的出现，UNet[4]已成为分割任务的基础 Backbone 网络。UNet以其对称的U形编码器-解码器架构和整体跳跃连接而闻名，这些连接有效地保留了关键的空间信息，并将编码器和解码器层中的特征合并，以解决复杂结构的分割问题。然而，在遥感领域，图像通常包含具有显著目标变化的大规模场景，基于卷积神经网络（CNN）[5]或Transformers[6]构建的UNet架构遇到了限制。

它们可能无法充分捕获全局上下文或表现出高计算复杂性，如图1（a）、（b）所示。因此，开发能够捕获全面的局部-全局信息的更高效架构至关重要。

近期的发展引入了创新的Mamba[7]架构，该架构擅长高效捕获全局上下文信息。Mamba专为长距离建模而设计，以其基于状态空间模型（SSM）[8]的计算效率而闻名。随后，Vision Mamba[9]和VMamba[10]将Mamba的架构扩展到计算机视觉领域，增强了Mamba的单向扫描机制。考虑到高效的全球上下文建模能力，Mamba架构非常适合于遥感图像处理，如图1（c）所示。PanMamba[11]，RMamba[12]，RS-Mamba[13]和RS3Mamba[14]探讨了将Mamba应用于处理遥感图像的方法。这些方法要么用Vision Mamba块替换网络并从头开始训练，要么直接应用预训练的Vision Mamba块。然而，它们很少考虑在遥感图像中整合局部和全局信息，这可能限制了它们充分利用预训练的CNN模型提供特征的能力。

在本文中，作者提出了CM-UNet，一个用于遥感图像语义分割的新框架。CM-UNet利用Mamba架构从CNN编码器聚合多尺度信息。它包括一个U形网络，具有提取多尺度文本信息的CNN编码器和一个具有设计的CSMamba块的解码器，用于高效的语义信息聚合。CSMamba块使用Mamba块以线性时间复杂度捕获长距离依赖性，并采用通道和空间注意力进行特征选择。

作为替代方法：

1. 作者提出了一个名为CM-UNet的基于Mamba的框架，以高效整合遥感图像语义分割中的局部-全局信息。
2. 作者设计了一个CSMamba块，将通道和空间注意力信息融入到Mamba块中，以提取全局上下文信息。此外，作者使用多尺度注意力聚合模块辅助跳跃连接，并采用多输出损失逐步监督语义分割。
3. 在三个著名的公开遥感数据集上进行的广泛实验，即ISPRS Potsdam、ISPRS Vaihingen和LoveDA，证明了所提出的CM-UNet的优越性。

II Methodology

作者的CM-UNet框架，如图2（a）所示，包含三个核心组件：基于CNN的编码器、MSAA模块和基于CSMamba的解码器。编码器采用ResNet来提取多级特征，而MSAA模块融合这些特征，取代了UNet的原始跳跃连接，增强了解码器的能力。

在CSMamba解码器中，CSMamba块的组合聚集了局部文本特征，以建立全面的语言理解。

CSMamba Block

大规模遥感语义分割受益于那些能够整合全局和局部信息处理能力的模型[15, 16]。最近的进展，包括采用自注意力机制的 Transformer 架构[17, 18]，在各类视觉任务中已经显示出显著的有效性。然而，这些模型经常遇到平方时间复杂度的限制，这对于处理大规模遥感图像的可扩展性和效率构成了挑战。

这强调了需要创新的方法，在保持 Transformer 优势的同时，减轻其计算需求。

其中DWConv表示深度卷积，CS指通道和空间注意力模块，2D-SSM是2D选择性扫描模块，而表示哈达玛积。原始的Mamba模型[7]通过顺序选择性扫描处理1-D数据，这适合NLP任务，但挑战非因果数据形式如图像。

遵循[10]，作者为图像语义分割引入了2D选择性扫描模块（2D-SSM）。

如图2.(c)所示，2D-SSM将图像特征展平为1D序列，并以四个方向扫描：从左上到右下，从右下到左上，从右上到左下，以及从左下到右上。这种方法通过选择性状态空间模型在每个方向捕获长距离依赖。

然后，合并方向序列以恢复2D结构。

Multi-Scale Attention Aggregation.

与此同时，通道聚合使用全局平均池化将维度降低到，然后通过卷积和ReLU激活生成通道注意力图。这张图被扩展以匹配输入的维度，并与空间细化图结合。因此，MSAA增强了后续网络层中的空间和通道特征。通过融入MSAA模块，得到的特征图用细化的空间和通道信息进行了丰富。

Multi-Output Supervision.

为了有效地监督解码器在逐步生成具有遥感图像的语义分割图，作者的CM-UNet架构在每个CSMamba块中融入了中间监督。这确保了网络的每个阶段都为最终的分割结果做出贡献，从而促进更精细和准确的输出。对于第个CSMamba块的中间输出是

III Experiments

Datasets

所提出的方法使用了ISPRS波茨坦、ISPRS法辛根以及LoveDA [19]遥感分割数据集进行评估。对于ISPRS波茨坦数据集，有14张图像用于测试，而剩余的23张图像（由于标注错误，排除了图像）用于训练目的。ISPRS法辛根数据集包括12个用于训练的图像块和4个用于测试的图像块。关于LoveDA数据集，训练集由1,156张图像组成，测试集由677张图像补充。作者采用平均F1分数（mF1）、平均交并比（mIoU）和总体准确率（OA）作为评估指标。

Implementation details

作者所有的实验都是在单个NVIDIA 3090 GPU上进行的，使用的是PyTorch框架。作者采用了AdamW优化器，基础学习率为6e-4，并采用余弦策略来调整学习率。遵循[15]，对于Vaihinge、Potsdam和LoveDA数据集，图像被随机裁剪成的块。在训练过程中，采用了如随机缩放（）、随机垂直翻转、随机水平翻转和随机旋转等增强技术，而训练周期设置为100，批量大小为16。在测试阶段，使用了测试时间增强（TTA），例如垂直翻转和水平翻转。

Performance Comparison

为了进行比较分析，作者纳入了一些值得注意的竞争方法作为基准，包括DeepLabV3+[20]，DANet[21]，ABCNet[22]，BANet[23]，CMTFNet[16]，UNetformer[15]，ESDINet[24]，BANet[23]和Segmenter[25]。这些方法使用了已确立的编码器架构，如R18[26]，VMamba[10]和Swin-Base[18]。

Iii-C1 The ISPRS Potsdam dataset

如表1所示，在ISPRS波茨坦测试集上，CM-UNet超越了其他竞争方法。它实现了93.05%的mF1、91.86%的OA和87.21%的mIoU，分别比UNetformer高出0.25%、0.56%和0.41%。值得注意的是，相比于传统方法如DANet和Segmenter，它在mIoU指标上分别提高了6.91%和6.51%，这突显了其预训练的ResNet主干网络和创新架构在空间特征学习方面的有效性。

与ESDINet、UNetformer和CMTFNet等近期模型相比，CM-UNet在多种指标上的优势进一步展示了其灵活性和有效性。图4中的定性比较进一步证明了其相对于UNetformer的优越性，尤其是在提取更清晰的建筑轮廓和减少错误分割方面。

Iii-C2 The ISPRS Vaihingen Dataset

表2展示了实验结果。CM-UNet实现了85.48%的mIoU，超过了竞争对手2.78%至16.08%。其mF1（92.01%）和OA（93.81%）也表现良好。在F1分数上，CM-UNet在多个类别中表现卓越，特别是在Imp.surf.、建筑物、Low.weg.和Car类别。它在Imp.surf.上比UNetFormer高出4.42%。这突显了Mamba捕捉不规则物体和在全球-局部关系中在RS中至关重要的能力。

如图5所示的可视化确认了其准确性，尤其是在辨别异常样本和细微变化如阴影方面。CM-UNet在imp. surf.和建筑物类别中的精确预测，强调了其感知全局长距离事实和空间上下文特征的能力。

Iv-C3 The LoveDA Dataset

表3展示了在LoveDA数据集上的结果。值得注意的是，作者的方法取得了性能，mIoU达到了52.17%。此外，CM-UNet在各种类别中表现突出，例如背景、建筑和道路。图6中的可视化强调了与UNetFormer相比，CM-UNet在描绘土地覆盖类别方面的优势。

它准确地捕捉到建筑边缘、道路和农业区域，即使在复杂城市场景和复杂的农业模式中也精确地接近 GT 标签。该方法在不同类别之间的一致性强调了其更高的分类准确性和增强的边缘检测能力，这对于精确的土地覆盖制图至关重要。这些结果证实了CM-UNet在大规模遥感图像上的有效性。

Further Analysis

Iv-D1 Effect of Model Architecture

消融研究对所设计模块的结果如表4所示。具体来说，仅采用多尺度注意力聚合（MSAA）模块就能带来改进，表明其在捕捉不同尺度上下文信息方面的有效性。

同样，融合多输出策略进一步增强了分割性能，证明了利用多个预测输出的好处。值得注意的是，MSAA和多输出模块的共同使用在所有指标上都带来了最高的增益，突显了这些架构组件在提高模型辨别复杂空间特征能力和整体分割准确度方面的协同效应。

Iv-D2 Model Complexity

表5展示了使用三种指标（浮点运算次数[FLOPs]，模型参数量和内存占用）比较模型复杂度的结果。值得注意的是，CM-UNet在这三个指标上实现了很好的平衡，具有更低的FLOPs和参数数量，同时内存占用也相对较小。

尽管如此，其mIoU结果却超过了其他模型，突显了其有利的性价比。

IV Conclusion

在本文中，作者介绍了CM-UNet，这是一个高效的框架，它利用了最近的Mamba架构进行遥感（RS）语义分割。作者的设计通过采用一种新型的UNet结构来应对大规模遥感图像中的显著目标变化。

编码器使用ResNet提取文本信息，而解码器则采用CSMamba块有效地捕获全局长距离依赖关系。此外，作者还集成了一个多尺度注意力聚合（MSAA）模块和多输出增强，以进一步支持多尺度特征学习。CM-UNet已经在三个遥感语义分割数据集上进行了验证，实验结果证明了作者方法的优势。

参考

[1].CM-UNet: Hybrid CNN-Mamba UNet for Remote Sensing Image Semantic Segmentation.