[논문 리뷰] Resource Optimized Neural Architecture Search for 3D Medical Image Segmentation

Woong Bae, Seungho Lee외 4명|arXiv (Cornell University)|2019. 09. 02.

Advanced Neural Network Applications참고 문헌 10인용 수 5

한 줄 요약

이 논문은 3D 의료 영상 분할을 위한 자원 최적화된 신경망 구조 탐색 프레임워크인 RONASMIS를 제안한다. 이 프레임워크는 단일 RTX 2080Ti(10.8GB GPU 메모리)를 사용하여 1GB 데이터셋에 대해 단지 1.39일 만에 최신 기술 수준의 성능을 달성한다. 매크로 검색과 파라미터 공유, 원소 합산을 통한 스킵 커넥션, 재학습을 방지함으로써, 데이터 증강, 앙상블, 테스트 시 시간 증강 없이도 수동으로 설계된 네트워크와 기존 NAS 방법을 모두 능가한다.

ABSTRACT

Neural Architecture Search (NAS), a framework which automates the task of designing neural networks, has recently been actively studied in the field of deep learning. However, there are only a few NAS methods suitable for 3D medical image segmentation. Medical 3D images are generally very large; thus it is difficult to apply previous NAS methods due to their GPU computational burden and long training time. We propose the resource-optimized neural architecture search method which can be applied to 3D medical segmentation tasks in a short training time (1.39 days for 1GB dataset) using a small amount of computation power (one RTX 2080Ti, 10.8GB GPU memory). Excellent performance can also be achieved without retraining(fine-tuning) which is essential in most NAS methods. These advantages can be achieved by using a reinforcement learning-based controller with parameter sharing and focusing on the optimal search space configuration of macro search rather than micro search. Our experiments demonstrate that the proposed NAS method outperforms manually designed networks with state-of-the-art performance in 3D medical image segmentation.

연구 동기 및 목표

• 기존 NAS 방법이 3D 의료 영상 분할에서 높은 계산 비용과 긴 학습 시간을 유발하는 문제를 해결하기 위해.
• 3D 의료 영상에서 데이터셋이 크고 이방성임을 감안해 GPU 메모리 사용량과 학습 시간을 줄이기 위해.
• 재학습이나 광범위한 데이터 증강 없이도 수동으로 설계된 아키텍처를 능가하는 NAS 프레임워크를 개발하기 위해.
• 제한된 계산 자원을 가진 임상 환경에서의 NAS 실용적 구현을 가능하게 하기 위해.
• 매크로 검색과 전략적인 아키텍처 선택이 3D 의료 분할에서 마이크로 검색을 능가할 수 있음을 보여주기 위해.

제안 방법

• 프레임워크는 파라미터 공유를 통해 효율적으로 매크로 검색 공간을 탐색하는 강화학습 기반 컨트롤러를 사용하여 학습 비용과 메모리 사용량을 감소시킨다.
• 검색 공간에는 입력 패치 크기, 다운샘플링 연산, 풀링 유형, 확장률, 스킵 커넥션 지점 등이 포함되며, 이는 최소한의 자원 소비로 최고의 성능을 달성하기 위해 선택된다.
• 파라미터와 GPU 메모리 소비를 줄이기 위해 스킵 커넥션에 원소 합산을 사용한다.
• 깊이 분할 컨볼루션을 표준 3D 컨볼루션으로 대체하여 메모리 효율성과 성능을 향상시킨다.
• 학습 중 자식 네트워크의 과적합을 방지하기 위해 드롭패스 정규화를 적용하고, 배치 정규화를 인스턴스 정규화로 교체하여 학습 안정성을 높였다.
• 컨트롤러는 150~500 에포크 동안 학습되고, 자식 네트워크는 에피소드당 3 에포크 동안 학습되며, 보상은 검증 세트의 Dice 스코어에 기반한다.

실험 결과

연구 질문

• RQ1 제한된 GPU 자원에서 효율적으로 작동하는 3D 의료 영상 분할을 위한 NAS 프레임워크를 설계할 수 있는가?
• RQ2 전략적인 아키텍처 선택이 있는 매크로 검색이 3D 의료 분할 과제에서 마이크로 검색을 능가하는가?
• RQ3 재학습, 데이터 증강, 앙상블 없이도 NAS 방법이 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ4 파라미터 공유와 원소 합산 스킵 커넥션은 성능을 유지하면서 메모리 사용량을 얼마나 효과적으로 줄이는가?
• RQ5 기존 모델이나 복잡한 데이터 파이프라인에 의존하지 않고도 컨트롤러 기반 NAS 접근 방식이 안정적이고 일관된 아키텍처 탐색을 달성할 수 있는가?

주요 결과

• RONASMIS는 심장 분할 과제에서 평균 Dice 스코어 92.72를 기록하여 nnUnet(92.70)과 SCNAS(89.99)를 능가하며 최신 기술 수준의 성능을 입증했다.
• 전립선 과제에서는 테스트 시 시간 증강이나 앙상블 없이도 Dice 스코어 75.71을 달성하여 nnUnet(74.54)과 SCNAS(65.30)를 초월했다.
• 단일 RTX 2080Ti(10.8GB GPU 메모리)를 사용하여 1GB 데이터셋에 대해 학습 시간을 1.39일로 단축시켰으며, 이는 이전 NAS 방법보다 크게 낮은 수준이다.
• 컨트롤러는 80% 이상의 시간 동안 동일하거나 매우 유사한 아키텍처를 제안하여 시간이 지남에 따라 엔트로피가 감소하고 보상이 증가함에 따라 안정적이고 일관된 탐색 행동을 보였다.
• 대부분의 최신 기술 수준 기준과 달리, 데이터 증강, 앙상블, 테스트 시 시간 증강, 후처리 없이도 뛰어난 성능을 달성했다.
• 프레임워크는 깊이 분할 컨볼루션을 표준 3D 컨볼루션으로 성공적으로 대체했으며, 인스턴스 정규화를 사용해 학습 안정성을 높여 3D 의료 영상에서 효율적이고 안정적인 NAS를 가능하게 했다.