[논문 리뷰] UNet++: Redesigning Skip Connections to Exploit Multiscale Features in Image Segmentation

Zongwei Zhou, Md Mahfuzur Rahman Siddiquee외 2명|arXiv (Cornell University)|2019. 12. 11.

Advanced Neural Network Applications참고 문헌 65인용 수 76

한 줄 요약

UNet++는 다층 규모의 특징을 융합하기 위해 촘촘하게 연결된 스킵 연결과 깊은 감독을 도입하여 U-네트의 자가 앙상블에서 다중 스케일 특징을 융합하고, 프루닝으로 추론 속도를 높여 시멘틱 및 인스턴스 의학 영상 분할을 향상시킨다.

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-넷을 학습시켜 깊은 감독을 통해 성능을 향상시킨다.
• 세그멘테이션 정확도에 큰 손실 없이 추론 속도를 높이기 위한 프루닝 메커니즘을 제공한다.
• 여러 백본 및 의학 영상 모달리티에 걸친 UNet++의 확장성을 입증한다.

제안 방법

• 인코더를 공유하지만 상호 얽힌 디코더를 갖는 서로 다른 깊이를 가진 U-넷의 앙상블을 구성한다.
• 촘촘한 스킵 연결로 디코더 특징을 대안 스케일의 인코더 특징과 밀도 있게 융합하도록 재설계한다( dense skip connections).
• 중간 디코더 노드에 깊은 감독을 적용하여 모든 구성 U-Nets를 동시에 학습시키고 추론 시 프루닝을 가능하게 한다.
• 픽셀 단위 교차 엔트로피와 소프트 Dice 손실을 결합한 하이브리드 손실을 정의하고, 등가 가중치로 디코더 전체에 걸쳐 합한다.
• 추론 시 더 깊은 디코더 경로를 제거하여 더 빠른 가지치기 버전(UNet++ Ld)을 생성하는 프루닝 체계를 도입한다.
• 향상된 인스턴스 분할을 위해 FPN 스킵 연결을 UNet++ 스타일 연결로 교체하는 Mask RCNN++을 입증한다.

실험 결과

연구 질문

• RQ1 UNet++가 다양한 의학 영상 데이터셋과 백본 백본에서 표준 U-Net보다 일관되게 시멘틱 분할 성능을 향상시키는가?
• RQ2 촘촘하게 연결된 스킵 연결과 깊은 감독이 다중 스케일 특징 융합과 수렴에 어떤 영향을 미치는가?
• RQ3 깊은 감독에 의한 프루닝으로 IoU/Dice의 큰 손실 없이 추론을 크게 가속화할 수 있는가?
• RQ4 UNet++가 다른 아키텍처(예: Mask R-CNN)와 모달리티(CT, MRI, EM 등)에 확장 가능한가?
• RQ5 가지치기된 서브 네트워크의 내장 학습이 성능 면에서 독립 학습보다 우수한가?

주요 결과

• UNet++는 6개 세그먼테이션 데이터셋과 여러 백본에서 일관되게 U-Net 및 Wide U-Net을 능가한다.
• 깊은 감독과 앙상블 유사 구조가 추가적인 IoU 이득을 가져오며, 여러 경우에 약 0.6~0.8 포인트의 향상을 보인다.
• Dense skip 연결은 다중 스케일 특징 융합을 효과적으로 가능하게 하며 핵 이미지, 뇌종양, 간, 폐결절 분할을 개선한다.
• Mask RCNN++가 UNet++ 설계로 FPN skip 연결을 대체하면 원래의 Mask R-CNN보다 핵 분할에서 인스턴스 분할 성능이 향상된다(IoU 93.28%에서 95.10%, Dice 87.91%에서 91.36%).
• Pruned UNet++(예: UNet++ L3)은 추론 속도 약 32.2% 증가, 메모리 75.6% 감소 등 상당한 속도 향상을 달성하며 IoU 저하도 약 0.6포인트로 modest하다.
• 프루닝된 서브-네트워크의 내장 학습은 고강도 프루닝 하에서도 독립 학습보다 우수한 성능을 보인다.