参数量减少40倍，推理速度提高6倍！UV-Mamba：结合变形卷积的网络如何克服SSM的内存问题？

极市平台

2024-09-12

↑ 点击蓝字关注极市平台

作者丨AI视界引擎

来源丨AI视界引擎

编辑丨极市平台

极市导读

这篇文章介绍了一种名为UV-Mamba的新型神经网络模型，该模型结合了变形卷积和状态空间模型，用于高分辨率遥感图像中精确检测城市村庄边界。 >>加入极市CV技术交流群，走在计算机视觉的最前沿

由于多样的地理环境、复杂的景观和高密度的居民区，利用遥感图像自动识别城中村边界是一个极其具有挑战性的任务。在本论文中，作者提出了一种新的、高效的神经网络模型UV-Mamba，用于在高分辨率遥感图像中准确检测边界。

UV-Mamba通过结合变形卷积（DCN）来抑制状态空间模型（SSM）中图像大小增加而带来的内存损失问题。其结构采用了一个编码器-解码器框架，包括一个拥有四个可变形状态空间扩展（DSSA）块的编码器用于高效的 multi-level语义提取，以及一个解码器用于集成提取的语义信息。

作者在北京和西安数据集上进行了实验，结果显示UV-Mamba达到了最先进的表现。具体而言，作者的模型在北京和西安数据集上的 IoU 分别达到了73.3%和78.1%，分别比现有最佳模型提高了1.2%和3.4%的 IoU，同时在推理速度上快6倍，参数数量上小40倍。源代码和预训练模型可在补充材料中找到。

I Introduction

城市村庄，作为城市化过程中的历史遗留物，由于其低层建筑和密集的建筑物，不理想的环保条件，以及过时的市政基础设施，在城市建设与经营管理中带来了较大的挑战。城市村庄的问题不仅关系到城市的形象美感和清洁，而且直接影响着居民的生活方式，公共安全和社会稳定。传统收集城市村庄信息的方法主要依赖于人工实地调查，这既耗时又费力 [8]。

为了实现城市村庄边界的自动识别，利用卫星影像进行图像分割技术的探索已引起了广泛关注。一些研究利用先进的语义分割模型，包括全卷积网络（FCN）和U-Net，来映射城市村庄区域 [13, 14, 15] 利用对抗学习来调整语义分割网络，以适应输入图像在不同领域的一致输出。UisNet [16] 通过空间-通道特征融合模块，结合遥感影像和建筑轮廓，增强分割准确度。UV-SAM [17] 利用普通模型和专用模型的优势，将SAM [18]的零样本学习能力应用于城市村庄边界识别任务。

然而，现有研究中准确界定城市村庄边界具有挑战性，主要原因有两点：

一是城市村庄的独特建筑特征，如高密度、狭窄的街道和多样化的建筑形式，使其具有固有的困难；

二是卷积神经网络（CNN）在捕捉全局信息方面存在局限性，以及 Transformer 的计算复杂性，如图1所示，进一步复杂化了这项任务。此外，当超高分辨率（UHR）遥感图像被划分为较小的小块时，空间特征和依赖关系可能会丢失。

为了解决上述问题，作者提出了UV-Mamba模型，该模型利用SSM的全局建模能力和线性复杂度的变形卷积的 spatial几何变形能力。作者的模型通过使用DCN来为感兴趣的区域分配更大的权重，从而改善SSM在长时间序列建模中的内存损失问题，从而提高SSM在不同序列上保留信息的能力。作者的架构的主要贡献如下：

作者引入了UV-Mamba，这是一种基于SSM的新颖而高效的建筑，它既保留了线性计算复杂性，又具有加强的全局建模能力。
作者设计了一种DSSA模块，它通过使用变形卷积为感兴趣的区域分配更大的权重，以减轻SSM在长距离建模过程中的记忆损失，从而提高SSM在扩大序列后保留信息的能力。
作者在中国两个城市北京和西安进行了广泛的实验，结果表明作者的方法取得了优越的性能，超过了基于CNN的现有方法和基于Transformer的方法。

在城市化过程中，城市村庄是一个不可忽视的现象。然而，准确地定义城市村庄的边界是一个具有挑战性的任务。城市村庄具有独特的建筑特征，如高密度、狭窄的街道和多样化的建筑形式，这使得界定其边界具有固有困难。此外，卷积神经网络（CNN）在捕捉全局信息方面的局限性，以及 Transformer 的计算复杂性，如图1所示，进一步复杂化了定义城市村庄边界的过程。因此，开发一种能够高效准确界定城市村庄边界的技术，对于城市规划和管理工作具有重要意义。针对这个挑战，作者提出了UV-Mamba模型，它利用SSM的全局建模能力和线性复杂度的变形卷积的 spatial几何变形能力。作者的模型通过使用DCN为感兴趣的区域分配更大的权重，以改善SSM在长时间序列建模中的内存损失问题，从而提高SSM在不同序列上保留信息的能力。

作者的架构的主要贡献如下：作者引入了UV-Mamba，这是一种基于SSM的新颖而高效的建筑，它既保留了线性计算复杂性，又具有加强的全局建模能力。

作者设计了一种DSSA模块，它通过使用变形卷积为感兴趣的区域分配更大的权重，以减轻

II Methodology

在本节，作者将阐述一种基于深度学习的图像分类方法。首先，作者将介绍数据集的预处理方法，然后将描述模型架构的选择和训练过程，最后给出模型在实际应用中的评估结果。

Preliminaries: State Space Model

状态空间模型是现代控制论中线性时不变系统的概念衍生出的一个概念。状态空间模型将一个维度的输入信号映射到一个维的潜在状态 , 然后将其投影到一个一维的输出信号。这个过程可以通过以下的线性一阶微分方程（ODE）来描述：

其中是状态转移矩阵, 和分别是投影矩阵。

为了更好地适应深度学习中如文本序列的离散输入，和使用零阶 hold ( ZOH 技术进行离散化, 引入可学习的时尺度参数 , 将连续状态空间模型转化为离散状态空间模型。离散化过程如下:

离散化后，第1式可表示为：

其中和分别表示和矩阵的离散版本。表示前一个状态信息, 表示当前状态信息。

模型概述

如图2 (a) 所示, 所提出的 UV-Mamba 模型由三个主要组成部分组成：一个具有可变卷积核大小的茎模块、一个层次化的多路径扫描编码器和一个轻量级的解码器。茎模块执行初始特征提取并下采样输入图像 by a factor of 2 , 由四个卷积层组成, 卷积核大小分别为和 ,填充分别为 3 和 1 , 步长分别为 2 和 1 。多路径扫描编码器由四个可变形状态空间扩展（DSSA）块组成, 每个阶段都可以将特征图大小减半, 从而相对于模型输入产生各种尺度的特征图: 。解码器包含四个上采样模块, 每个模块通过转置卷积将特征图从编码器中上采样两倍, 随后由两个卷积进行特征融合。最后, 双线性插值用于将图像恢复到输入大小。

Deformable State Space Augmentation Block

对于UHR远程 sensing密集城市环境，两个主要挑战是：提高像素 Level 的表示和确保SSM的鲁棒全局建模以进行准确的边界提取。为了解决这些挑战，作者设计了一个DSSA块，如图2（b）所示，它包括以下部分：分块嵌入（patch embeddings）、空间适应可变形增强器（SADE）、多路径扫描SSM模块（MSSM）和分块合并。特别的是，作者的SADE和MSSM模块作为中间模块堆叠了两次。通过SADE对感兴趣区域赋予更重的权重，可以缓解全局建模过程中由于SSM导致的内存损失。这种方法在获得线性复杂度的同时，增强了SSM模型的全局建模能力，使其可以更有效地对建筑物进行区分，如图3所示。

多路径扫描SSM模块（MSSM）。 一系列研究【21, 22, 23, 24】已经表明，在基于SSM的模型中，增加扫描方向的数量对于实现全面的全球建模能力至关重要。为了更好地划分城乡界限，作者聚合了来自八个方向的扫描结果（水平、垂直、对角和反对角，前后都可以），以捕捉周围结构的复杂空间关系，并对上下文环境有全面的了解。为了更好地适应不同的输入大小，作者引入了Mix-FFN，它比传统的位置编码（positional encoding）更有效【25】来提供位置信息，通过在Feed-Forward网络中应用3x3卷积。

空间适应可变形增强器（SADE）。 如图2（c）所示，SADE的设计采用了类似于 Transformer 的结构。【29】。通过利用可变形卷积的空间几何变形学习能力，它可以更好地适应城市村庄的多样化空间分布特征。具体来说，作者使用DCNv4操作符对空间特征进行增强，因为其速度快、计算效率高。这个过程如下：

其中表示聚合组的总数。对于第组, 表示与位置无关的投影权重, 是第个采样点的调制 scalar, 表示分块输入特征图, 是网格采样位置的偏移的。然后, 作者使用 Mix-FFN对提取的特征进行聚合, 这可以减少计算复杂度, 同时保持模型的表示能力。

III Experiments

Experimental Settings

数据集：作者使用来自北京和西安的两个具有独特建筑风格的中国城市[17]的数据集，这两个城市由于其显著的地理位置差异而具有不同的建筑方式。这两个城市都有传统和现代建筑的混搭，形成了复杂的都市结构，这对作者的模型在提取都市村庄边界方面提出了挑战。北京数据集包含531张图像，西安数据集包括205张图像。作者将这些数据集划分为训练、验证和测试集，比例为6:2:2。每张图像的分辨率均为1024X1024，以确保包含主要的城市信息。

实现细节：作者的实验在单张Tesla V100 GPU上进行，训练100个周期。为了防止过拟合并提高泛化能力，作者在所有实验中应用了统一的数据增强策略，其中包括随机旋转、水平翻转和垂直翻转。模型在城市景观数据集[31]上预训练，然后在对都市村庄数据集进行微调。在预训练过程中，作者使用Adam[32]优化器，初始学习率为0.001。学习率在第一十个周期内 Warm up ，然后逐渐减小到1e-6。跨熵损失[33]在预训练阶段用于优化模型的性能。

预训练权重随后在对都市村庄数据集上进行微调。对于北京和西安数据集上的微调，作者继续使用Adam优化器。学习率在第一个三十个周期内 Warm up ，然后逐渐减小到1e-6。具体来说，对于北京数据集，作者将学习率设置为0.0004，并使用Dice损失函数[34]。对于西安数据集，学习率设置为0.0002，并采用交叉熵损失函数。模型的准确性采用交点与一致性（IoU）、准确率（ACC）和总体准确性（OA）进行评估。效率通过对参数（Params, M）和浮点运算次数（Flops, G）进行评估，分别以#P和#F表示，以便于阅读表格。

Ablation Studies

图像大小： 为了评估上下文信息和空间特征对城市村庄边界检测的影响，作者使用不同大小的输入图像来评估模型性能，结果如表1所示。实验结果显示，随着图像大小的增加，城市村庄检测的准确性持续提高，这可能是由于这些区域的空间持续分布。这一发现强调了使用UHR遥感图像进行精确边界检测的重要性。

DSSA模块： 为了评估UV-Mamba中DSSA模块的有效性，作者在表2中展示了不同模型变体在北京和西安数据集上的分割性能。结果表明，在去掉SADE模块后，模型的性能降低了2.4%和5.5%；同样，去掉MSSM模块后，性能降低了2.8%和6.7%。这些结果强调了准确城市村庄分割对强大的全局建模能力的重要性。此外，作者在DSSA模块内尝试了SADE和MSSM模块的各种位置组合。结果表明，当将SADE和MSSM模块并排时，性能次优，分别达到了72.7%和74.9%的IoU。相反，将SADE模块放在MSSM模块后面会导致整体模型性能最差，这表明SSM的长序列建模限制导致特征图信息丢失，从而误导模型。总之，这些结果表明SADE可以部分补充SSM的全局建模能力，帮助在SSM模型中处理高分辨率遥感图像时减轻内存损失问题。

Comparison to the State-of-the-Arts

如图III所示，UV-Mamba优于先进的市区村庄识别模型 [38, 35, 36, 37, 38]，在两个数据集上都达到了最新的性能。可视化的分割结果见图4。关于分割精度，与先前最佳的城市村边界识别模型UV-SAM相比，作者的模型在两个数据集上的IoU提高了1%-3%，而参数大小只有UV-SAM的1/40。此外，在ACC和OA的准确性指标中也观察到了类似的性能提升。