企业官网建设流程全解析

Slim UNETR：面向计算资源受限场景的混合 Transformer 轻量化三维医学图像分割网络（TMI 2024）

论文代码链接：

GitHub - aigzhusmart/Slim-UNETR: An Efficient, High-Quality 3D Segmentation for Medical Image Analysis with Constrained Computational Resources · GitHubhttps://github.com/aigzhusmart/Slim-UNETR

Abstract：

Hybrid transformer‑based segmentation approaches have shown great promise in medical image analysis. However, they typically require considerable computational power and resources during both training and inference stages, posing a challenge for resource‑limited medical applications common in the field. To address this issue, we present an innovative framework called Slim UNETR, designed to achieve a balance between accuracy and efficiency by leveraging the advantages of both convolutional neural networks and transformers. Our method features the Slim UNETR Block as a core component, which effectively enables information exchange through self‑attention mechanism decomposition and cost‑effective representation aggregation. Additionally, we utilize the throughput metric as an efficiency indicator to provide feedback on model resource consumption. Our experiments demonstrate that Slim UNETR outperforms state‑of‑the‑art models in terms of accuracy, model size, and efficiency when deployed on resource‑constrained devices. Remarkably, Slim UNETR achieves 92.44% dice‑accuracy on BraTS2021 while being 34.6× smaller and 13.4× faster during inference compared to Swin UNETR.

翻译：

基于混合 Transformer 的分割方法在医学图像分析领域展现出巨大潜力。然而，这类方法在训练与推理阶段均需要消耗巨大的算力与硬件资源，这对于医疗领域普遍存在的资源受限应用场景构成巨大挑战。针对该问题，本文提出创新网络框架 Slim UNETR，该网络结合卷积神经网络与 Transformer 两者的优势，实现分割精度与运行效率之间的平衡。该方法的核心是 Slim UNETR 模块，通过对自注意力机制进行分解、低成本特征表征聚合，实现高效的特征信息交互。除此之外，本文将吞吐量（throughput）作为效率评价指标，用来反映模型的资源消耗情况。实验结果表明，在资源受限设备上部署时，Slim UNETR 在分割精度、模型参数量、推理效率上均优于现有 SOTA 方法。值得注意的是，在 BraTS2021 数据集上 Slim UNETR 取得 92.44% 的 Dice 精度；对比 Swin UNETR，模型参数量仅为其 1/34.6，推理速度提升 13.4 倍。

创新点

（1）提出面向资源受限场景的轻量化混合 U 型三维医学分割框架Slim UNETR。

融合卷积神经网络与 Transformer 各自优势，在保证分割精度的同时大幅压缩模型参数量、提升推理速度，解决现有 CNN‑Transformer 混合模型算力开销大、难以在低算力医疗设备部署的痛点。

（2）设计核心Slim UNETR Block 模块。

内部构建上下文集成流程，拆解传统自注意力机制；由局部特征聚合 LRC、LineConv、全局稀疏 Transformer GST、局部反向扩散 LRD 四部分组成。通过稀疏 token 策略，避免对全部 token 执行注意力运算，以较低计算成本完成局部与全局特征之间的信息交互。

（3）提出将吞吐量（throughput）作为模型实际效率评价指标。

指出传统 FLOPs 无法反映内存访问开销、硬件并行特性，吞吐量可以更加真实地衡量资源受限硬件上模型的实际推理性能，为轻量化医学分割模型提供更贴合落地场景的评估方式。

（4）在 BraTS2021、MSD2019 脑肿瘤、MSD2021 肝血管三套大规模三维医学数据集开展充分实验。

Slim UNETR 的整体网络框架

该模型采用带有跳跃连接的 U 型编码器-解码器架构，用于生成多尺度分层特征图。

网络左侧为编码器，一共包含 4 个阶段（Stage1~Stage4），每一个阶段均由 3D 深度卷积 DW Conv 与 Slim UNETR Block 顺序堆叠而成，输入是维度为 H×W×D×4 的四维脑肿瘤多模态体数据；随着层级逐级下采样，特征图的空间尺寸不断缩小，依次变为 H/4×W/4×D/4 、 H/8×W/8×D/8、H/16×W/16×D/16、H/32×W/32×D/32，通道数同步提升，在这一过程中网络逐层提取由细到粗的多尺度上下文特征。

图右下角放大框展示了核心的 Slim UNETR Block，也就是每个阶段内完成特征交互的基础单元，模块内部依次执行局部特征聚合 LRC、线性卷积 LC、全局稀疏 Transformer GST、局部反向扩散 LRD，最终采用残差连接，输入特征与模块输出相加，完成局部细节与全局长距离依赖的低成本信息交互，规避传统完整自注意力带来的二次复杂度开销。

网络右侧为解码器分支，采用 3D 转置卷积 TS Conv 搭配 Block 完成逐级上采样，逐步恢复特征图的空间分辨率；编码器每一级输出的特征图，会通过跳跃连接直接传递至解码器对应的层级，与上采样后的特征逐元素相加，弥补下采样过程丢失的细节信息。在完成全部解码阶段之后，特征送入分割头 Seg Head，最终输出 H×W×D×3 的分割预测结果，对应脑肿瘤的三类分割子区域。

整体架构继承 UNETR 系列 U 型编解码的范式，但依靠轻量化 Slim UNETR Block，在构建分层特征的同时大幅降低了混合 Transformer 架构的计算开销，适配资源受限设备的三维医学图像分割任务。

Slim UNETR Block 的内部结构（核心！）

该模块构建了一个低成本瓶颈结构，内部设计上下文集成流程，实现特征在局部表征与全局表征之间的转换，以分解自注意力的方式，用较低计算开销完成三维医学体数据的全局 - 局部特征交互。模块输入特征为X_in，整体是多段残差串联的结构，每一个子模块的输出都和输入特征进行逐元素相加，也就是残差连接，有效缓解深层网络训练时的梯度消失问题。

模块首先进入LRC 局部表征聚合（Local Representations Congregation）分支，对应图左侧放大子模块。LRC 模块由 BatchNorm、逐点 3D 卷积 Pointwise Conv3D、3D 深度卷积 Depthwise Conv3D、BatchNorm 以及另一层逐点卷积串联构成。其前向计算原文公式写为：

LRC 依靠深度卷积捕捉三维空间内局部邻域细节，逐点卷积完成通道维度的特征融合，专门负责提取医学图像的局部纹理、边界等细粒度信息；计算完成后使用残差相加，保留原始输入特征，避免局部信息丢失。

经过 LRC 之后，特征送入LineConv，LineConv 采用 Pointwise Conv3D、GELU 激活、Pointwise Conv3D 的结构，完成通道间非线性变换，对应公式：

LineConv 在不破坏三维空间结构的前提下做通道映射，对局部聚合后的特征做进一步平滑与表征增强。

接下来特征经过 Patch Partition 操作，划分特征块后送入GST 全局稀疏 Transformer（Global Sparse Transformer），也就是图右下放大的子模块，这是整个 Block 最核心的创新部分，用于替代传统完整自注意力。GST 首先使用步幅为r的逐点卷积，对特征做稀疏采样，得到 Q、K、V 三组特征，仅在稀疏采样后的 token 上执行多头自注意力，而不是对全部 token 做全局计算，以此降低 O(N^2)的二次复杂度。完成多头自注意力计算后，再通过转置卷积恢复到原始分辨率，最后搭配逐点卷积完成特征映射。GST 的计算过程原文表示为：

GST 只在少量 token 之间建模长距离依赖，兼顾全局上下文建模能力，同时大幅削减计算量与参数量，解决传统 Transformer 在三维体图像上开销巨大的问题。

GST 输出的特征送入LRD 局部反向扩散（Local Reverse Diffusion）模块。LRD 的作用是将 GST 学到的全局上下文信息扩散回原始稠密特征网格，把全局语义信息重新注入到每个局部空间位置，弥补稀疏采样造成的局部细节丢失，计算公式：

补充：反向扩散

GST 阶段用步幅 r 的逐点卷积，把一整块窗口内很多个体素压缩为 1 个中心 token；在 LRD 中使用相同步幅 r 的 3D 转置卷积做反向映射，该卷积核会把这一个中心 token 存储的全局特征值，按照卷积核权重，重新分配、加权叠加回它原本对应的整个窗口所有体素位置上，相当于把单个 token 携带的全局信息，扩散给到窗口里面每一个空间坐标，完成信息的 “辐射分发”（分发、辐射本质就是转置卷积的加权上采样运算），将特征图分辨率还原为送入 GST 之前的稠密尺寸；完成上采样之后，再将得到的特征张量和 GST 输入的原始稠密特征做残差相加，原始特征中 LRC 提取的局部纹理信息得以完整保留，不会被覆盖，最终每一个体素既保留原有的局部细节，又融合了从稀疏中心 token 传递过来的全局上下文，以此弥补 GST 稀疏采样仅对少量 token 做注意力造成的空间信息缺失，同时保证张量维度匹配后续网络的计算要求。

LRD 完成全局信息向局部空间的回传，实现全局表征到局部表征的转换。经过 LRD 之后再次接入一层 LineConv，做最后的通道非线性特征校准，得到模块最终输出特征X_out，完整模块整体前向过程：

整体来看，Slim UNETR Block 这套上下文集成流程，完成了“局部提取→通道映射→稀疏全局建模→全局信息回传”的闭环，在同一个模块内部交替完成局部特征与全局特征之间的转换。和 UNETR、Swin UNETR 中直接对全部 token 做自注意力的方式不同，Slim UNETR Block 通过分解自注意力、稀疏 token 采样，在保留全局上下文建模能力的前提下显著降低计算开销，是整个 Slim UNETR 轻量化、高精度的核心基础单元，该模块可以堆叠在编码器的每一个 Stage 中，分层提取多尺度三维医学特征。

补充：Slim‑UNETR 对自注意力机制的"分解"原理

传统完整的多头自注意力，需要全部 token 两两之间计算相关性，复杂度是O(N^2)，三维医学体数据 token 数量巨大，直接计算显存和算力开销极高。Slim‑UNETR 并不是修改注意力数学公式，而是把完整的 “局部提取‑全局注意力‑全局信息回写” 这一个整体自注意力流程拆解成三个子阶段，用卷积与稀疏注意力分工协作，以此实现对自注意力机制的分解。

第一步由LRC 局部表征聚合先用卷积在各个局部窗口内完成局部信息的聚合浓缩，把稠密的原始 token 集合提炼出一批代表局部窗口信息的中心 token；不再拿全部原始 token 进入注意力计算。

第二步送入GST 全局稀疏 Transformer，只对上面得到的少量中心代表 token 执行标准多头自注意力，完成跨窗口的全局依赖建模，这一步只在稀疏采样后的少量 token 之间做(Q,K,V)计算，避开全部 token 两两交互带来的平方复杂度。

第三步依靠LRD 局部反向扩散，把 GST 在少量中心 token 上学到的全局语义，通过转置卷积反向扩散，把全局信息重新分发回原始全部稠密空间位置。

简单来说，它把原本一步完成的 “全部 token 同时做全局自注意力”，拆解成：卷积做局部浓缩 → 少量代表 token 做稀疏全局注意力 → 再把全局信息扩散回全部位置。用卷积承接局部运算，仅用少量 token 承担全局注意力运算，实现分解。

Slim UNETR Block 内部的上下文集成流程

按照论文原文描述，该流程旨在提升大规模分割模型性能的同时降低计算开销，整体由 LRC 局部表征聚合、GST 全局稀疏 Transformer、LRD 局部反向扩散三个核心子过程构成，完成局部特征提取→全局上下文建模→全局信息回灌到局部的完整闭环。

图 (a) 为局部表征聚合 LRC（Local Representations Congregation）。该模块将输入三维体特征切分为多个 3D 窗口嵌入块，在每个窗口内部，利用卷积对窗口内的局部 Patch 做特征聚合，提取图像局部纹理、边界这类细粒度信息，把原始稠密特征浓缩为局部集中特征。该步骤将原始输入拆分成更小的子集，降低后续模块的处理难度，专门负责捕捉三维医学图像上的邻域局部信息。

图 (b) 是全局稀疏 Transformer GST（Global Sparse Transformer）。它接收 LRC 输出的局部集中特征，选取中心注意力 Patch 作为计算基准，仅在稀疏采样得到的少量 Patch 之间执行自注意力计算，学习跨局部窗口的高层全局依赖。不同于传统 Transformer 对全部 token 做注意力，GST 只对稀疏选取的 Patch 建模长距离关联，以此规避原生自注意力 O(N^2) 的高计算复杂度，在大幅减少计算量的前提下学习全局语义模式。

图 (c) 为局部反向扩散 LRD（Local Reverse Diffusion）。该模块的功能是把 GST 学到的全局语义知识，重新扩散、融合回每一个局部表征当中。它将全局注意力得到的特征，反向映射回原始稠密的 3D 窗口网格，生成局部扩散特征，把全局上下文信息注入各个局部位置，补偿稀疏注意力采样过程丢失的空间细节，完成从全局表征到局部表征的转换。

模型效率评估

论文指出传统 FLOPs 仅代表理论计算量，无法反映内存读写、硬件并行等实际硬件约束，不能代表低算力设备上真实推理性能。因此论文引入吞吐量 Throughput作为核心效率评价指标，单位为 vol/s，代表每秒能够处理的三维体数量。

吞吐量直接在目标硬件上实测得到，可以真实体现模型在资源受限设备上的推理速度。作者同时结合参数量、Dice 指标做综合评估。实验结果显示，Slim UNETR 相比 Swin UNETR，在 Dice 更高的情况下，参数量缩小 34.6 倍，吞吐量提升 13.4 倍，证明该轻量化方案可以在不损失精度的前提下大幅降低部署成本。

消融实验

论文设计多组消融实验，验证上下文集成流程中 LRC、GST、LRD 各个组件的有效性，同时探究稀疏比例、窗口大小等超参对性能与效率的影响。

实验逐步移除 LRC、GST、LRD 模块，观察分割 Dice 指标的下降幅度。实验结果证明，移除任意一个子模块，分割精度都会明显下降，证明 LRC 局部特征聚合、GST 稀疏全局注意力、LRD 局部反向扩散三者缺一不可，协同完成局部与全局特征交互。除此之外，作者针对 GST 的稀疏采样比例做对照，权衡稀疏度带来的精度 - 速度取舍；同时测试不同窗口尺寸，分析窗口大小对全局建模能力与显存占用的影响，选取最优配置。消融实验完整验证了 Slim UNETR Block 每一个设计的必要性。