分享自:

Slim UNETRV2:用于资源受限医疗便携设备的3D图像分割

期刊:IEEE Transactions on Medical ImagingDOI:10.1109/tmi.2025.3602145

Slim UNETRv2: 面向资源受限医疗便携设备的3D医学图像分割模型

一、研究概述

本研究报告介绍了一项发表于 IEEE Transactions on Medical Imaging(Volume 45, Issue 2, February 2026)的原创性研究。该研究由中国科学院深圳先进技术研究院的 Yan Pang、Ying Hu 和 Qiong Wang(通讯作者),华南理工大学的 Jiaming Liang,广州大学的 Junming Yan,以及香港科技大学的 Hao Chen 共同完成。研究团队提出了名为 Slim UNETRv2 的轻量级三维医学图像分割框架,专门为资源受限的医疗便携设备设计,以实现高精度、高速度和低推理抖动(Inference Jitter)的分割性能。

二、学术背景与研究目标

随着数字治疗和辅助诊断技术的快速发展,医疗便携设备对图像分割模型提出了日益严苛的要求。这些设备必须具备三个核心特征:紧凑的设计以便于在不同环境中运输和使用;即时响应能力以应对紧急场景或持续监测;以及在各种条件下的一致运行稳定性,这要求分割模型具备低推理抖动。

当前主流的医学图像分割方法主要采用两种技术路线:基于注意力机制(Attention-based)的方法(如 Transformer 和 Mamba 模型)擅长捕获长距离依赖和上下文信息,但计算复杂度高、内存消耗大;基于卷积(Convolutional)的方法则以其轻量级设计、高效的参数管理和空间归纳偏置著称,但受限于局部感受野,难以捕获全局特征。现代混合框架试图整合两者优势,在编码器和解码器中同时或选择性使用混合技术,然而这种策略增加了架构复杂性,导致推理速度减慢、推理抖动增加,这对于需要快速可靠处理的实时应用构成了挑战。

面对上述问题,本研究的目标是开发一种创新的设备端分割框架——Slim UNETRv2,在继承前代模型 Slim UNETR 成功经验的基础上,通过简化架构设计和引入新颖的特征聚合机制,实现分割精度、推理速度和稳定性之间的最优平衡。

三、研究方法与技术路线

(一)整体框架设计

Slim UNETRv2 采用四阶段的U形编码器-解码器架构,输入端处理尺寸为 H×W×D×C 的三维医学图像。编码器各阶段通过深度卷积(Depthwise Convolution)和实例归一化(Instance Normalization)层将特征图逐步下采样至原始尺寸的二分之一、四分之一、八分之一和十六分之一。解码器则利用转置卷积(Transposed Convolution)将特征图上采样回原始尺寸,最后通过分割头输出分割概率。与以往在编码器和解码器内部同时采用混合技术的模型不同,Slim UNETRv2 的关键创新在于仅在编码器和解码器中使用基础卷积操作,从而最大化执行速度并最小化推理抖动,同时在每个层级的跳跃连接(Skip Connections)中引入专门设计的 Slim UNETRv2 Block 来弥补卷积操作在全局特征捕获方面的不足。

(二)空间融合模块(Spatial Fusion Module, SFM)

在编码器的每个阶段,输入特征首先经过实例归一化处理以稳定数据分布,随后进行深度卷积以降低计算复杂度并优化特征表示。SFM 模块由三个并行分支构成:第一个分支应用 3×3 深度卷积、实例归一化和激活函数;第二个分支采用逐点卷积(Pointwise Convolution)与实例归一化及自适应激活函数;两个分支的输出融合后,通过第三个残差连接分支与原始输入相加,以保持数据完整性并增强特征多样性,最后经多层感知机(MLP, Multi-Layer Perceptron)进行非线性变换。值得注意的是,SFM 采用了自适应激活策略:在编码器的浅层阶段(第1和第2阶段)使用高斯误差线性单元(GELU, Gaussian Error Linear Unit)以利于浅层表示的有效激活,而在深层阶段则切换至 Swish 激活函数以维持深层网络中的信息传播和梯度流动。这种设计显著提升了特征处理的整体有效性和效率。

(三)Slim UNETRv2 Block 与稀疏三阶段选择性状态空间和自注意力机制

Slim UNETRv2 Block 是框架的核心组件,集成在编码器和解码器之间的每个层级跳跃连接中。在每个阶段,输入的更新表示 x(维度为 h×w×d×c)首先经过下采样处理,采样率 r 根据阶段不同进行调整:第1和第2阶段的 r 设为4,后续阶段 r 降至2,这种设计在初始阶段更为激进地降低分辨率以优化计算资源,而在深层阶段保留更多空间信息。下采样后的特征图 x_sq 尺寸变为 h/r × w/r × d/r × c。

该模块采用一种名为稀疏三阶段选择性状态空间和自注意力机制(Sparse Tri-SSA)的创新架构,由三个并行分支组成:前向分支(Forward Branch)利用状态空间模型(State-Space Model)沿一个方向顺序处理输入,以线性计算复杂度高效捕获长距离依赖;反向分支(Reverse Branch)沿相反方向扫描输入,获取逆序上下文信息;层间分支(Inter-slice Branch)通过聚合相邻切片的特征来建模跨层关系,且限定在局部邻域内进行稀疏计算以减少内存和计算开销。三个分支的输出经求和后重塑为原始下采样尺寸,随后进行自注意力(Self-Attention)计算,使每个元素能够关注所有其他元素,从而全面理解复杂的空间关系。最后,注意力矩阵经重塑和转置卷积上采样恢复至原始输入尺寸,得到最终的特征配置。

(四)损失函数与训练配置

Slim UNETRv2 采用结合软骰子损失(Soft Dice Loss)和焦点损失(Focal Loss)的复合损失函数。软骰子损失准确度量预测输出与真实标签之间的重叠度,焦点损失则针对难以分类的样本进行重点学习,以缓解数据类别不平衡问题。

研究使用了三个公开数据集进行验证:BraTS 2021 数据集包含 1251 例三维 MRI 样本,具有 T1、T1GD、T2、T2-FLAIR 四种模态,标注了三类肿瘤亚区;BraTS 2023 数据集同样包含 1251 例三维脑部 MRI 数据,专注于脑膜瘤病例;MM-WHS 数据集包含 120 例多模态全心脏图像,其中 60 例为心脏 CT 图像,60 例为 MRI 图像,涵盖七个心脏亚结构。模型在 MONAI 平台上开发,使用四个 NVIDIA RTX 4090 GPU 进行训练。训练过程中采用了随机轴翻转、强度缩放与偏移、旋转、弹性形变、伽马校正等多种数据增强策略。

四、主要实验结果

(一)分割精度对比

在 BraTS 2021 数据集上,Slim UNETRv2 取得了 93.89% 的平均骰子相似系数(DSC)和 2.90 mm 的 95% 豪斯多夫距离(HD95),在所有对比模型中表现最优。相较于 LightM-UNet(DSC 91.52%,HD95 4.34 mm),DSC 提升了 2.37 个百分点,HD95 降低了 1.44 mm;相较于 NNMamba(DSC 91.03%,HD95 5.33 mm),DSC 提升了 2.86 个百分点,HD95 降低了 2.43 mm;与其前代模型 Slim UNETR(DSC 92.44%,HD95 5.36 mm)相比,DSC 提升了 1.45 个百分点,HD95 降低了 2.46 mm。

在 BraTS 2023 数据集上,Slim UNETRv2 取得了平均 DSC 94.18% 的优异表现,其中全肿瘤(WT)为 95.25%,增强肿瘤(ET)为 91.88%,肿瘤核心(TC)为 95.41%,较 SegMamba-V2(91.60%)提高了 2.58 个百分点,较 LightM-UNet(88.69%)提高了 5.49 个百分点;平均 HD95 为 2.80 mm,相较于 SegMamba-V2 降低 0.43 mm。

在 MM-WHS 心脏分割数据集上,Slim UNETRv2 在 CT 模态下取得了平均 DSC 92.45% 的最佳成绩,尤其在右心室(RVBC)的分割上达到 96.22%,较第二名的 UNETR++(91.71%)高出 4.51 个百分点;在 MRI 模态下取得平均 DSC 88.32%,右心室分割 DSC 为 93.92%,较 UNETR++(91.67%)提升 2.25 个百分点。

(二)执行速度与推理抖动

在 NVIDIA RTX 4090 上,Slim UNETRv2 实现了 140.856 FPS 的吞吐量,是 LightM-UNet(4.958 FPS)的约 28 倍,是 NNMamba(3.596 FPS)的近 39 倍,是 Slim UNETR(65.548 FPS)的 2.14 倍。在推理抖动方面,Slim UNETRv2 仅为 0.225 毫秒,而 LightM-UNet 为 8.838 毫秒(约 39 倍差异),NNMamba 为 8.777 毫秒,Slim UNETR 为 0.497 毫秒(约 2.2 倍差异),充分展现了其出色的稳定性和可预测性。

(三)跨平台性能验证

研究在 22 种不同芯片平台上进行了速度测试,涵盖 8 款桌面 GPU、4 款笔记本 GPU 和 10 款移动芯片。在桌面端,Slim UNETRv2 在 RTX 4090 上吞吐量达 140.856 FPS,在 RTX 3090 上为 95.253 FPS,分别是 SegMamba 的约 18 倍和 Slim UNETR 的 2.18 倍。在笔记本端,于 GTX1650 上达到 32.451 FPS,约为 Swin UNETR(3.212 FPS)的 10 倍。在移动芯片上,以 Snapdragon 855 Plus 为例,Slim UNETRv2 的推理时间为 356 毫秒,约为 NNMamba(2743 毫秒)的 7.7 倍快;在 Dimensity 8200 上推理时间为 330 毫秒,较 SegMamba(28350 毫秒)减少了约 86 倍。值得注意的是,SegMamba 在 Snapdragon 855 Plus 上因内存限制运行失败,而 Slim UNETRv2 展现出更强的资源受限环境适应能力。

(四)消融实验

消融实验验证了各组件对性能的贡献。在 Sparse Tri-SSA 机制的配置中,SM-SL(三方向Mamba结合自注意力)配置在保持合理参数量(21.59M)和计算量(31.91 GFLOPs)的前提下,实现了最佳的分割性能。框架设置方面,基线模型(无任何增强模块)仅取得 86.36% 的 DSC,仅添加 Slim UNETRv2 Block 提升至 92.87%,同时集成 Block 和 SFM 后达到最优的 95.26%(针对该特定实验设置),验证了两者在平衡计算效率和分割精度方面的协同作用。下采样率消融则确定了 4-2-2-2 配置为帕累托最优解,在保证 93.89% DSC 的同时维持了 140.85 FPS 的高吞吐量。

五、研究结论与应用价值

本研究成功开发了 Slim UNETRv2,一个资源高效的 3D 医学图像分割模型。通过在编码器和解码器中统一使用基础卷积操作,并将高效的稀疏 Tri-SSA 机制置于跳跃连接中的 Slim UNETRv2 Block,该模型实现了高分割精度、快速推理和极低推理抖动的三方面最佳权衡。在三个公开数据集上的实验结果一致表明,Slim UNETRv2 在分割精度上超越了当前最先进的模型,同时在执行速度和稳定性上取得数量级提升。

该研究的科学价值在于提出了一种新的架构设计范式:通过将复杂的全局建模能力从主干网络转移到跳跃连接中,在简化主路径以提升速度的同时,保证了对全局上下文的有效捕获。其应用价值则体现在为医疗便携设备上的实时医学图像分析提供了实用化解决方案,使得在救护车、远程患者监护等资源受限场景下实现高精度即时影像诊断成为可能。

六、研究亮点

架构创新:首次提出在编码器和解码器中仅使用基础卷积操作,将混合注意力机制专门集成于跳跃连接的设计方案,有效解耦了速度优化与精度保障的需求。

Sparse Tri-SSA 机制:独创性地融合前向状态空间建模、逆向上下文捕获和跨层关系聚合三种策略,并以稀疏方式进行层间计算,实现了对局部结构特征和全局上下文信息的高效协同捕获,同时将计算开销控制在极低水平。

全面的跨平台验证:在涵盖桌面、笔记本和移动端共 22 种芯片上的广泛测试,充分证明了模型在多样化资源受限环境下的鲁棒性和实用性。

突破性性能指标:在 BraTS 2021 数据集上同时取得了 93.89% 的最高 DSC 和 140.856 FPS 的最高吞吐量,且推理抖动仅 0.225 毫秒,为实时临床部署提供了强有力支持。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com