【魔改Mamba系列】UNetMamba:一种用于高分辨率遥感图像语义分割的高效UNet类Mamba模型

科技   2024-11-19 10:05   中国香港  

点击下方深度学习爱好者”,选择加p"星标"或“置顶

论文信息

题目:UNetMamba: An Efficient UNet-Like Mamba for Semantic Segmentation of High-Resolution Remote Sensing Images

UNetMamba:一种用于高分辨率遥感图像语义分割的高效UNet类Mamba模型

作者:Enze Zhu, Zhan Chen, Dingkai Wang, Hanru Shi, Xiaoxuan Liu, and Lei Wang

源码:https://github.com/EnzeZhu2001/UNetMamba

论文创新点

  1. Mamba分割解码器(MSD)的设计:作者提出了一个基于Mamba的分割解码器MSD,这是一个即插即用的解码器,用于从多尺度特征图中高效解码语义信息。MSD通过将VMamba的视觉状态空间(VSS)块转移到解码侧,显著减少了参数数量,同时利用长距离建模能力在全局感受野中准确解码复杂信息。
  2. 局部监督模块(LSM)的引入:为了解决MSD在解码过程中可能忽略局部语义信息的问题,作者设计了一个基于CNN的局部监督模块LSM,以增强对局部细节的感知。LSM通过两种不同尺度的卷积和一个辅助损失函数增强局部语义信息的感知,并通过训练专用设计节省推理成本,确保应用效率。
  3. UNetMamba模型的整体架构:作者提出了一个UNet类模型UNetMamba,它结合了预训练的ResT作为编码器和提出的MSD作为解码器,以及LSM来增强局部语义信息的感知。该模型不仅在两个知名的遥感数据集上实现了最先进的性能,而且具有轻量级和低成本的优势。

摘要

高分辨率遥感图像的语义分割对于土地覆盖映射、城市规划和灾害评估等下游应用至关重要。现有的基于Transformer的方法在准确性和效率之间存在约束,而最近提出的Mamba因其高效性而闻名。因此,为了克服这一困境,我们提出了UNetMamba,一种基于Mamba的UNet类语义分割模型。它结合了一个Mamba分割解码器(MSD),可以高效地解码高分辨率图像中的复杂信息,以及一个局部监督模块(LSM),该模块仅在训练时使用,但可以显著增强局部内容的感知能力。广泛的实验表明,UNetMamba在LoveDA上mIoU提高了0.87%,在ISPRS Vaihingen上提高了0.39%,同时通过轻量级设计、较小的内存占用和降低的计算成本实现了高效率。

关键字

遥感、语义分割、Mamba
UNet类模型。它具有轻量级和低成本的优势,同时也在两个知名的高分辨率遥感图像数据集上实现了最先进的(SOTA)性能。

方法

如图1(a)所示,UNetMamba的整体架构是在U形框架下构建的,选择了预训练的ResT作为编码器,而解码器是提出的MSD。此外,LSM被设计用来增强局部语义信息的感知。

ResT编码器

如图1(a)所示,ResT编码器以Efficient Transformer Block(ETB)为核心,由四个阶段组成,用于捕获多尺度特征图。为了避免原始多头自注意力的二次计算成本,ETB采用了高效的多头自注意力(EMSA):
其中引入了实例归一化IN()和1×1卷积Conv()来提高效率。当应用于高分辨率遥感图像时,EMSA可以确保我们UNetMamba的高编码效率。

Mamba分割解码器

由于高分辨率遥感图像在展平时可以被视为非常长的序列,与Transformer相比,Mamba在线性扩展能力上具有自然优势。实际上,在先驱研究中,使用Mamba进行语义分割的好处已经得到了确立。然而,使用预训练的VMamba作为编码器的模型(>30M)往往过于沉重,无法实现我们的高效率目标。因此,我们将VMamba的基本单元,即VSS块,仅纳入解码侧。

VSS块的结构

VSS块的结构如图1(b)所示。经过patch扩展的2D特征图F首先经过层归一化,然后分成两个流。主流进一步线性嵌入为F':
随后,F'经过3×3深度卷积,然后SiLU激活为F'':
2D选择性扫描(SS2D)模块随后在四个不同方向上扫描F'',在全局感受野和线性复杂度下解码语义信息:
其中扩展和合并操作对应于VMamba,而S6操作表示Mamba的选择性扫描空间状态序列模型。然后,经过层归一化和与旁路流的逐元素乘法,旁路流也经历了相同的线性嵌入和Silu激活。最后,VSS块通过带有残差连接的线性层输出FM_SD。基于上述操作,MSD在四个不同阶段解码多尺度特征图,最终通过1×1卷积头输出语义分割结果。

局部监督模块

VSS块中SS2D的大感受野往往是双刃剑,导致MSD在解码时部分忽略了局部语义信息。然而,这些细节在高分辨率图像中是进一步提高语义分割准确性的关键因素。因此,提出了基于CNN的灵活LSM来增强局部细节的感知,如图1(c)所示。考虑到高分辨率场景中土地覆盖的相对较小规模,我们在LSM中采用了两个并行卷积分支,核大小分别设置为3和1。每个分支后面跟着一个批量归一化层和ReLU6激活:
其中i ∈ {1, 3}表示核大小,而FM_SD表示相应阶段MSD的输出特征。两个分支然后合并为。最终,通过一个dropout层,一个与MSD中的分割头一致的1×1卷积,以及一个上采样层到原始图像大小,LSM获得了
在训练期间,我们在解码器的第2-4阶段添加了具有不同输入通道的LSM块,然后通过求和获得分割结果,该结果进一步用于计算辅助损失函数。

损失函数

训练UNetMamba所使用的损失函数包括两部分,一部分是用于整体优化的主损失Lp,另一部分是用于局部监督的辅助损失La。为了确保UNetMamba的整体分割性能,选择了两个经典损失函数共同构成主损失Lp:
其中 分别代表Dice损失和交叉熵损失。N表示样本数量,K表示类别数量。 是样本n的真实标签的第k个元素的one-hot编码,而 $ y(n)k L{\text{ce}} L_a L_p L_a $ 的效果,然后实现最佳性能,训练UNetMamba的整体损失函数被制定为加权方程:
其中权重因子 设置为0.4。

实验结果

声明

本文内容为论文学习收获分享,受限于知识能力,本文对原文的理解可能存在偏差,最终内容以原论文为准。本文信息旨在传播和学术交流,其内容由作者负责,不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题,请及时与我们联系,我们将在第一时间回复并处理。
下载1:Pytoch常用函数手册
在「深度学习爱好者」公众号后台回复:Pytorch常用函数手册,即可下载全网第一份Pytorch常用函数手册,涵盖Tensors介绍、基础函数介绍、数据处理函数、优化函数、CUDA编程、多线程处理等十四章章内容。
下载2:Python视觉实战项目52讲
小白学视觉公众号后台回复:Python视觉实战项目即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目,助力快速学校计算机视觉。

交流群

欢迎加入公众号读者群一起和同行交流,目前有SLAM、三维视觉驶、计算摄影、检测、分割、识别、医学影像、GAN算法竞赛等微信群(以后会逐渐细分),请扫描下面微信号加群,备注:”昵称+学校/公司+研究方向“,例如:”张三 + 上海交大 + 视觉SLAM“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告,否则会请出群,谢谢理解~

深度学习爱好者
分享机器学习、深度学习和Python等知识与原理,每天分享深度学习与计算机视觉领域的经典和最新的论文总结,带读者一起跟踪前言科技!
 最新文章