[논문 리뷰] DoDNet: Learning to segment multi-organ and tumors from multiple partially labeled datasets

Jianpeng Zhang, Yutong Xie외 2명|arXiv (Cornell University)|2020. 11. 20.

Advanced Neural Network Applications참고 문헌 37인용 수 8

한 줄 요약

DoDNet는 공유 인코더-디코더 아키텍처와 작업 조건부로 적응하는 세그먼테이션 헤드를 갖춘 동적 온디맨드 네트워크로, 부분적으로 레이블이 붙은 CT 데이터셋에서 여러 복부 기관과 종양을 세그먼테이션한다. 이는 MOTS 데이터셋에서 사전 훈련한 후에 다운스트림의 레이블 제한된 작업에 대해 강력한 전이 학습을 가능하게 하며, 일곱 가지 세그먼테이션 작업에서 최신 기술 수준(SOTA) 성능을 달성한다.

ABSTRACT

Due to the intensive cost of labor and expertise in annotating 3D medical images at a voxel level, most benchmark datasets are equipped with the annotations of only one type of organs and/or tumors, resulting in the so-called partially labeling issue. To address this, we propose a dynamic on-demand network (DoDNet) that learns to segment multiple organs and tumors on partially labeled datasets. DoDNet consists of a shared encoder-decoder architecture, a task encoding module, a controller for generating dynamic convolution filters, and a single but dynamic segmentation head. The information of the current segmentation task is encoded as a task-aware prior to tell the model what the task is expected to solve. Different from existing approaches which fix kernels after training, the kernels in dynamic head are generated adaptively by the controller, conditioned on both input image and assigned task. Thus, DoDNet is able to segment multiple organs and tumors, as done by multiple networks or a multi-head network, in a much efficient and flexible manner. We have created a large-scale partially labeled dataset, termed MOTS, and demonstrated the superior performance of our DoDNet over other competitors on seven organ and tumor segmentation tasks. We also transferred the weights pre-trained on MOTS to a downstream multi-organ segmentation task and achieved state-of-the-art performance. This study provides a general 3D medical image segmentation model that has been pre-trained on a large-scale partially labelled dataset and can be extended (after fine-tuning) to downstream volumetric medical data segmentation tasks. The dataset and code areavailableat: https://git.io/DoDNet

연구 동기 및 목표

• 높은 레이블링 비용으로 인해 부분적으로 레이블이 붙은 데이터셋만 이용 가능한 3D 의료 영상에서 다기관 및 종양 세그먼테이션 문제 해결.
• 기존의 다중 네트워크 및 다중 헤드 아키텍처가 부분적으로 레이블이 붙은 데이터에 대해 비효율적이고 유연성이 떨어지는 문제 극복.
• 작업 및 입력 컨텍스트에 따라 적응적으로 세그먼테이션 필터를 생성할 수 있는 단일 효율적이고 영리한 모델 개발.
• 다운스트림 세그먼테이션 작업을 위한 사전 훈련 및 전이 학습을 지원하기 위해 부분적으로 레이블이 붙은 대규모 데이터셋(MOTS) 구축.
• DoDNet의 일반화 능력과 전이 가능성 실증: 실제 레이블 제한된 의료 영상 시나리오에 적용 가능

제안 방법

• 작업별로 실시간으로 필터를 생성하는 단일 동적 세그먼테이션 헤드를 갖춘 공유 인코더-디코더 아키텍처 제안.
• 동적 헤드의 필터 생성을 조절하기 위해 작업별 사전 지식을 통합하는 작업 인코딩 모듈 도입.
• 입력 이미지 특징과 작업 임베딩을 기반으로 동적 컨볼루션 커널을 생성하는 경량 제어 네트워크 설계.
• 각 데이터셋이 오직 한 종류의 기관 또는 종양 유형에 대한 레이블만 포함하는 다수의 부분적으로 레이블이 붙은 데이터셋에서 모델을 종합적으로 훈련.
• 각 배치에서 활성 헤드만 업데이트하는 전략을 사용한 다중 작업 학습 방식으로 계산 자원 낭비 최소화.
• 다운스트림 작업에 대한 미세조정 전에 920건의 부분적으로 레이블이 붙은 복부 CT 케이스로 구성된 대규모 데이터셋인 MOTS에서 DoDNet 사전 훈련 수행.

실험 결과

연구 질문

• RQ1 단일 딥러닝 모델이 다수의 부분적으로 레이블이 붙은 데이터셋에서 여러 기관과 종양을 효과적으로 세그먼테이션할 수 있는가?
• RQ2 동적이고 작업 조건부인 세그먼테이션 헤드가 고정 또는 다중 헤드 아키텍처에 비해 효율성과 성능을 어떻게 향상시키는가?
• RQ3 MOTS와 같은 대규모 부분적으로 레이블이 붙은 데이터셋에서 사전 훈련한 모델이 레이블이 제한된 다운스트림 작업으로 얼마나 잘 일반화되는가?
• RQ4 제안된 동적 헤드 메커니즘이 다중 네트워크 또는 다중 헤드 대비 더 빠른 추론 속도와 더 낮은 파라미터 효율성을 제공하는가?
• RQ5 DoDNet가 완전히 레이블링되지 않은 데이터가 필요한 조건에서도 다양한 기관 및 종양 세그먼테이션 벤치마크에서 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• DoDNet는 MOTS 벤치마크에서 일곱 가지 기관 및 종양 세그먼테이션 작업 전반에서 최신 기술 수준(SOTA) 성능 달성하며, 정확도와 추론 속도 모두에서 기존 방법을 능가.
• 모델의 동적 헤드 덕분에 다중 헤드 네트워크보다 더 빠른 추론 속도 확보했으며, 반복적 추론 시 계산 오버헤드가 극히 미미.
• MOTS 데이터셋에서의 사전 훈련은 BCV 다기관 세그먼테이션 작업에 대한 미세조정 시 수렴 속도 향상과 최종 성능 향상에 크게 기여.
• BCV 테스트 세트에서 MOTS 사전 훈련을 적용한 DoDNet는 평균 Dice 스코어 86.44% 달성하여, 초기 학습(85.30%) 및 다른 SOTA 방법보다 표면 거리 및 하우스도르프 거리 지표에서 뛰어난 성능 확보.
• 사전 훈련된 DoDNet는 평균 표면 거리(MSD)를 1.46에서 1.17로 감소시키고, 하우스도르프 거리(HD)를 19.67에서 15.62로 감소시켜 비교된 모든 방법 중 최고의 MSD와 두 번째로 높은 Dice 스코어 기록.
• MOTS 사전 훈련된 가중치를 사용한 전이 학습은 BCV 챌린지에서 성능 향상에 상당한 기여를 하여, 미리 보지 못한 레이블 제한된 작업에 대한 강력한 일반화 능력 입증.