[논문 리뷰] nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation

Fabian Isensee, Jens Petersen외 9명|arXiv (Cornell University)|2018. 09. 27.

Advanced Neural Network Applications인용 수 390

한 줄 요약

nnU-Net은 일반 U-Net를 기반으로 한 자체 적응형의 완전 자동 파이프라인으로, 아키텍처, 전처리, 학습, 추론을 데이터셋 특성에 맞춰 조정하고, 수작업 튜닝 없이 Medical Segmentation Decathlon 데이터셋에서 최상의 성능을 달성합니다.

ABSTRACT

The U-Net was presented in 2015. With its straight-forward and successful architecture it quickly evolved to a commonly used benchmark in medical image segmentation. The adaptation of the U-Net to novel problems, however, comprises several degrees of freedom regarding the exact architecture, preprocessing, training and inference. These choices are not independent of each other and substantially impact the overall performance. The present paper introduces the nnU-Net ('no-new-Net'), which refers to a robust and self-adapting framework on the basis of 2D and 3D vanilla U-Nets. We argue the strong case for taking away superfluous bells and whistles of many proposed network designs and instead focus on the remaining aspects that make out the performance and generalizability of a method. We evaluate the nnU-Net in the context of the Medical Segmentation Decathlon challenge, which measures segmentation performance in ten disciplines comprising distinct entities, image modalities, image geometries and dataset sizes, with no manual adjustments between datasets allowed. At the time of manuscript submission, nnU-Net achieves the highest mean dice scores across all classes and seven phase 1 tasks (except class 1 in BrainTumour) in the online leaderboard of the challenge.

연구 동기 및 목표

• 비아키텍처적이고 자동으로 구성된 구성 요소들이 다양한 의료 분할 작업에서 특화된 아키텍처보다 우수한 성능을 낼 수 있음을 보인다.
• U-Net 토폴로지, 전처리, 학습, 추론 및 후처리에 대해 각 데이터셋에 맞춰 완전히 자동으로 파이프라인을 적응시키는 것을 도입한다.
• Medical Segmentation Decathlon에서 프레임워크를 평가하고 최첨단 벤치마크와 비교한다.

제안 방법

• 원본 U-Net에서의 작은 아키텍처 변화만으로도 3가지 U-Net 변형(2D U-Net, 3D U-Net, 및 U-Net Cascade) 은행을 사용한다.
• 각 데이터셋의 기하학에 맞춰 입력 패치 크기, 메모리 사용량, 풀링 연산을 자동으로 조정한다.
• 완전 자동화된 파이프라인을 정의한다: 전처리(자르기, 재샘플링, 정규화), 학습(손실: dice+CE, Adam 옵티마이저, 학습률 스케줄, 교차 검증), 추론(오버랩이 있는 패치 기반 타일링 및 테스트 시간 증강), 후처리(가장 큰 연결 요소).
• 적용 가능 시 데이터 증강(회전, 스케일링, 탄력 변형, 감마 보정, 대칭)을 사용하고 cascade의 경우 단계적 학습을 수행한다.
• 데이터 세트당 여러 모델을 앙상블하고(두 개 중 다섯 개 중 세 개를 선정하는 방식 포함) 교차 검증 평균 전경 Dice 점수를 기반으로 제출에 가장 잘 수행하는 모델을 선택한다.

실험 결과

연구 질문

• RQ1 다양한 의료 영상 데이터셋에서 태스크 특수 조정 없이도 일반화가 높은, 일반 U-Net을 이용한 자기 적응 프레임워크가 달성될 수 있나?
• RQ2 자동화된 전처리, 학습, 추론 및 후처리 구성 요소가 아키텍처 혁신에 비해 분할 성능에 어떤 영향을 미치는가?
• RQ3 교차 데이터셋 일반화에 대한 모델 앙상블 및 cascade 전략의 영향은 어떻게 나타나는가?

주요 결과

• nnU-Net 프레임워크는 제출 시점의 Medical Segmentation Decathlon 리더보드에서 7개 phase 1 작업의 평균 Dice 점수에서 최고를 기록했다(BrainTumour 클래스 1의 한 가지 예외를 제외하고).
• 데이터셋의 기하학(입력 패치 크기, 풀링 토폴로지)에 자동으로 적응함으로써 서로 다른 크기와 모달리티의 이미징에 대해 네트워크가 효과적으로 작동하도록 한다.
• 큰 이미지 크기에 대처하고 전체 해상도에서의 정제를 가능하게 하는 cascade 접근 방식(3D U-Net cascade)이 필요에 따라 적용된다.
• 구체적인 학습 프로토콜 세부사항(손실 = dice + 교차 엔트로피, Adam 옵티마이저, 학습률 스케줄링)과 광범위한 데이터 증강)이 강건성과 성능에 상당한 기여를 한다.
• 추론은 중첩이 있는 패치 기반 타일링과 테스트 시간 증강을 사용하여 각 보셀당 최대 64개의 예측을 집계해 강건성을 높인다.
• 자동 앙상블 전략(두 모델 앙상블)과 교차 검증 지표에서 가장 잘 수행하는 모델의 선택이 최종 제출 성능을 좌우한다.