[论文解读] UNet#: A UNet-like Redesigning Skip Connections for Medical Image Segmentation

Ledan Qian, Xiaofang Zhou+2 more|arXiv (Cornell University)|May 24, 2022

Radiomics and Machine Learning in Medical Imaging被引用 9

一句话总结

该论文提出UNet#——一种新型UNet类架构，通过结合密集连接与全尺度跳跃连接，增强多尺度特征之间的特征聚合，从而提升器官和病灶边界分割性能。通过引入深度监督实现模型剪枝，并采用分类引导模块减少误报，UNet#在多个2D和3D医学图像分割基准上达到最先进性能。

ABSTRACT

As an essential prerequisite for developing a medical intelligent assistant system, medical image segmentation has received extensive research and concentration from the neural network community. A series of UNet-like networks with encoder-decoder architecture has achieved extraordinary success, in which UNet2+ and UNet3+ redesign skip connections, respectively proposing dense skip connection and full-scale skip connection and dramatically improving compared with UNet in medical image segmentation. However, UNet2+ lacks sufficient information explored from the full scale, which will affect the learning of organs' location and boundary. Although UNet3+ can obtain the full-scale aggregation feature map, owing to the small number of neurons in the structure, it does not satisfy the segmentation of tiny objects when the number of samples is small. This paper proposes a novel network structure combining dense skip connections and full-scale skip connections, named UNet-sharp (UNet\#) for its shape similar to symbol \#. The proposed UNet\# can aggregate feature maps of different scales in the decoder sub-network and capture fine-grained details and coarse-grained semantics from the full scale, which benefits learning the exact location and accurately segmenting the boundary of organs or lesions. We perform deep supervision for model pruning to speed up testing and make it possible for the model to run on mobile devices; furthermore, designing two classification-guided modules to reduce false positives achieves more accurate segmentation results. Various experiments of semantic segmentation and instance segmentation on different modalities (EM, CT, MRI) and dimensions (2D, 3D) datasets, including the nuclei, brain tumor, liver, and lung, demonstrate that the proposed method outperforms state-of-the-art models.

研究动机与目标

• 解决UNet2+在捕捉全尺度上下文信息方面的局限性，以及UNet3+在小样本对象上训练数据有限时的性能瓶颈。
• 设计一种统一架构，同时利用密集连接与全尺度跳跃连接，增强语义层级间的特征融合。
• 通过基于深度监督的模型剪枝实现高效推理，适用于移动设备部署。
• 利用分类引导模块减少误报，提升边界分割精度。

提出的方法

• 提出一种新型UNet#架构，采用#形结构，结合密集跳跃连接与全尺度跳跃连接，实现多尺度特征聚合。
• 在八个解码器分支中实施深度监督，包括最终输出和四个中间层，以支持分层表征学习与模型剪枝。
• 通过基于权重大小的剪枝方法对深度监督分支进行剪枝，在保持性能的同时显著减少模型参数量。
• 在不同阶段集成两个分类引导模块（CGM），通过优化决策边界，抑制过度分割与误报。
• 采用多尺度监督，利用不同分辨率的真值标签，指导解码器各层级的特征学习。
• 使用交叉熵损失与Dice损失联合训练分割头与分类头，以同时优化定位与类别判别能力。

实验结果

研究问题

• RQ1 结合密集连接与全尺度连接的混合跳跃连接设计，是否能提升医学图像分割中的特征表示能力？
• RQ2 在多个解码器分支中实施深度监督，是否能有效实现模型剪枝而不引起性能下降？
• RQ3 分类引导模块是否能减少误报并提升器官与病灶分割的边界精度？
• RQ4 UNet#在多种医学成像模态与维度下，与最先进模型相比表现如何？

主要发现

• 在BraTs19数据集上，UNet#取得95.60%的Dice分数与91.68%的IoU，优于所有对比的SOTA模型。
• 引入分类引导模块（CGM）后，UNet#在BraTs19上达到96.14%的Dice分数与92.08%的IoU，较基础模型分别提升0.54%与0.4%。
• 在Lits17数据集上，UNet#取得97.29%的Dice与94.79%的IoU，超越所有基线模型，包括TransUNet与AttnUNet。
• CGM模块在Lits17上进一步将Dice分数提升0.95%，IoU提升0.52%，证实其在减少误报方面的有效性。
• 通过深度监督实现的模型剪枝使参数量最多减少97倍（从9.71M降至0.1M），同时保持强大的推理精度。
• 剪枝后的L¹模型（0.1M参数）在DSB2018上实现88.5%的IoU，证明其在移动设备部署的可行性。