[论文解读] UNet++: Redesigning Skip Connections to Exploit Multiscale Features in Image Segmentation

Zongwei Zhou, Md Mahfuzur Rahman Siddiquee+2 more|arXiv (Cornell University)|Dec 11, 2019

Advanced Neural Network Applications参考文献 65被引用 76

一句话总结

UNet++ 引入密集连接的跳跃连接和深度监督，以从内置的 U-Nets 集成中融合多尺度特征，提升语义分割和实例分割的医疗影像分割性能，同时通过修剪实现更快的推断。

ABSTRACT

The state-of-the-art models for medical image segmentation are variants of U-Net and fully convolutional networks (FCN). Despite their success, these models have two limitations: (1) their optimal depth is apriori unknown, requiring extensive architecture search or inefficient ensemble of models of varying depths; and (2) their skip connections impose an unnecessarily restrictive fusion scheme, forcing aggregation only at the same-scale feature maps of the encoder and decoder sub-networks. To overcome these two limitations, we propose UNet++, a new neural architecture for semantic and instance segmentation, by (1) alleviating the unknown network depth with an efficient ensemble of U-Nets of varying depths, which partially share an encoder and co-learn simultaneously using deep supervision; (2) redesigning skip connections to aggregate features of varying semantic scales at the decoder sub-networks, leading to a highly flexible feature fusion scheme; and (3) devising a pruning scheme to accelerate the inference speed of UNet++. We have evaluated UNet++ using six different medical image segmentation datasets, covering multiple imaging modalities such as computed tomography (CT), magnetic resonance imaging (MRI), and electron microscopy (EM), and demonstrating that (1) UNet++ consistently outperforms the baseline models for the task of semantic segmentation across different datasets and backbone architectures; (2) UNet++ enhances segmentation quality of varying-size objects -- an improvement over the fixed-depth U-Net; (3) Mask RCNN++ (Mask R-CNN with UNet++ design) outperforms the original Mask R-CNN for the task of instance segmentation; and (4) pruned UNet++ models achieve significant speedup while showing only modest performance degradation. Our implementation and pre-trained models are available at https://github.com/MrGiovanni/UNetPlusPlus.

研究动机与目标

• 通过解决分割模型中未知的最佳深度来推动对标准 U-Net 的改进。
• 开发一种灵活的跳跃连接方案，在编码器-解码器路径之间聚合多尺度特征。
• 在单一架构内实现多深度 U-Nets 的训练，以通过深度监督提升性能。
• 提供一种修剪机制以在推断中加速，同时不出现显著的分割准确度损失。
• 展示 UNet++ 在多种骨干网络和医疗影像模态上的可扩展性。

提出的方法

• 形成一个由不同深度的 U-Net 组成的集成，它们共享一个编码器但解码器相互交织。
• 重新设计跳跃连接，将解码器特征与不同尺度的编码器特征密集融合（密集跳跃连接）。
• 对中间解码节点应用深度监督，以同时训练所有组成的 U-Net，并在推断阶段实现修剪。
• 定义将像素级交叉熵和 Soft Dice 损失相结合的混合损失，在各解码器之间等权求和。
• 引入一种修剪方案，在推断阶段可以丢弃更深的解码路径，创建更快的修剪变体（UNet++ Ld）。
• 通过用 UNet++ 风格的连接替换 FPN 跳跃连接，演示 Mask RCNN++ 以改善实例分割。

实验结果

研究问题

• RQ1 在多组医疗影像数据集和多种骨干网络上，UNet++ 是否始终优于 U-Net 的语义分割性能？
• RQ2 密集连接的跳跃连接和深度监督如何影响多尺度特征融合与收敛？
• RQ3 通过深度监督的修剪能否在显著提速推断的同时保持最小的 IoU/Dice 损失？
• RQ4 UNet++ 是否可扩展到其他架构（如 Mask R-CNN）和模态（CT、MRI、EM 等）？
• RQ5 嵌入式训练的修剪子网络在提升较浅网络性能方面是否优于独立训练？

主要发现

• UNet++ 在六个分割数据集和多种骨干网络上持续优于 U-Net 与宽 U-Net。
• 深度监督结合类集成架构在若干情形下带来额外的 IoU 增益，大约在 0.6–0.8 点左右。
• 密集跳跃连接实现有效的多尺度特征融合，提升细胞核、脑肿瘤、肝脏和肺结节分割。
• Mask RCNN++ 采用 UNet++ 设计在核分割的实例分割上优于原始 Mask R-CNN（IoU 从 93.28% 提升到 95.10%，Dice 从 87.91% 提升到 91.36%）。
• 修剪后的 UNet++（如 UNet++ L3）实现显著的加速（推理大约快 32.2%，内存减少 75.6%），IoU 下降较小（约 0.6 点）。
• 嵌入式训练的修剪子网络在尤其在激进修剪情况下，优于孤立训练。