[논문 리뷰] Automated Model Design and Benchmarking of 3D Deep Learning Models for COVID-19 Detection with Chest CT Scans

Xin He, Shihao Wang외 7명|arXiv (Cornell University)|2021. 01. 14.

COVID-19 diagnosis using AI참고 문헌 40인용 수 8

한 줄 요약

이 논문은 흉부 CT 영상에서 코로나19 검출을 위한 효율적인 3D 딥러닝 모델을 자동으로 설계하기 위해 Gumbel Softmax를 활용한 미분 가능 신경망 아키텍처 탐색(DNAS) 프레임워크를 제안한다. 결과적으로 얻어진 CovidNet3D 모델들은 인간이 설계한 기준 모델들보다 높은 정확도(예: COVID-CTset에서 94.27%)를 달성하면서도 최대 41배 작아졌으며, CAM 시각화를 통해 질병 관련 영역을 강조함으로써 해석 가능성을 확인하였다.

ABSTRACT

The COVID-19 pandemic has spread globally for several months. Because its transmissibility and high pathogenicity seriously threaten people's lives, it is crucial to accurately and quickly detect COVID-19 infection. Many recent studies have shown that deep learning (DL) based solutions can help detect COVID-19 based on chest CT scans. However, most existing work focuses on 2D datasets, which may result in low quality models as the real CT scans are 3D images. Besides, the reported results span a broad spectrum on different datasets with a relatively unfair comparison. In this paper, we first use three state-of-the-art 3D models (ResNet3D101, DenseNet3D121, and MC3\_18) to establish the baseline performance on the three publicly available chest CT scan datasets. Then we propose a differentiable neural architecture search (DNAS) framework to automatically search for the 3D DL models for 3D chest CT scans classification with the Gumbel Softmax technique to improve the searching efficiency. We further exploit the Class Activation Mapping (CAM) technique on our models to provide the interpretability of the results. The experimental results show that our automatically searched models (CovidNet3D) outperform the baseline human-designed models on the three datasets with tens of times smaller model size and higher accuracy. Furthermore, the results also verify that CAM can be well applied in CovidNet3D for COVID-19 datasets to provide interpretability for medical diagnosis.

연구 동기 및 목표

• 흉부 CT 영상에서 코로나19 검출을 위한 3D 딥러닝 모델에 대한 공정하고 일관된 벤치마킹의 부족을 해결하기 위해.
• 실제 3D CT 데이터에서 2D 모델의 한계를 극복하고 3D 모델을 개발하고 평가하기 위해.
• 다양한 신경망 아키텍처 탐색(DNAS)을 활용해 고성능이면서도 컴팩트한 3D CNN 아키텍처를 자동으로 설계하기 위해.
• 예측 결과의 해석 가능성을 높이기 위해 클래스 활성화 맵(CAM)을 통합하여 임상적 신뢰도를 향상시키기 위해.
• 세 가지 공개 데이터셋에서 최신 3D 모델들(ResNet3D101, DenseNet3D121, MC3_18)을 기반으로 강력한 기준 모델을 수립하기 위해.

제안 방법

• 세 가지 공개 흉부 CT 데이터셋(Clean-CC-CCII, MosMedData, COVID-CTset)에서 최적의 3D CNN 아키텍처를 탐색하기 위해 다이퍼런셜 신경망 아키텍처 탐색(DNAS) 프레임워크를 활용한다.
• Gumbel Softmax 기법을 사용하여 단일 슈퍼넷의 엔드 투 엔드 학습을 가능하게 하여, 한 번에 하나의 서브넷을 효율적으로 최적화할 수 있도록 한다.
• 각 레이어에 대해 다양한 후보 연산(예: 컨볼루션 유형, 커널 크기 등)을 포함하는 슈퍼넷을 정의함으로써 탐색 공간을 설정한다.
• 최종 모델인 CovidNet3D는 정확도, 정밀도, 재현율, F1 점수 및 모델 크기와 같은 표준 지표를 사용해 훈련 및 평가된다.
• 클래스 활성화 맵(CAM)을 적용하여 주의 맵을 시각화하고, 모델이 실제 임상적으로 관련 있는 영역(예: 지상유리 농탁)에 집중하고 있음을 검증한다.
• 데이터 전처리 과정으로 NIfTI 파일을 PNG로 변환하여 호환성을 확보하지만, 이 과정에서 정보 손실가능성이 존재한다.

실험 결과

연구 질문

• RQ1 Gumbel Softmax를 활용한 DNAS 프레임워크는 흉부 CT 영상에서 코로나19 검출을 위한 컴팩트하고 정확한 3D CNN 아키텍처를 효과적으로 탐색할 수 있는가?
• RQ2 자동으로 탐색된 CovidNet3D 모델은 인간이 설계한 3D 모델들(ResNet3D101, DenseNet3D121, MC3_18)과 비교해 여러 데이터셋에서 정확도와 모델 크기 측면에서 어떻게 성능을 내는가?
• RQ3 클래스 활성화 맵(CAM) 통합을 통해 영상에서 질병 관련 영역을 강조함으로써 임상적으로 의미 있는 해석 가능성을 제공하는가?
• RQ4 왜 일부 모델은 MosMedData 데이터셋에서 다른 두 데이터셋보다 성능이 열 劣한가? 그리고 NIfTI에서 PNG로의 형식 변환은 성능에 어떤 영향을 미치는가?
• RQ5 랜덤 시드 선택이 NAS 기반 모델의 훈련 및 성능에 어느 정도 영향을 미치는가?

주요 결과

• CovidNet3D-L 모델은 모든 세 가지 데이터셋에서 가장 높은 정확도를 기록하며, 모든 기준 모델을 초월한다.
• CovidNet3D-S(8.36 MB)는 COVID-CTset 데이터셋에서 94.27%의 정확도를 달성하였으며, 이는 ResNet3D101(325.21 MB)보다 0.4% 높은 성능이며, 크기가 41배 작다.
• DNAS 프레임워크는 다양한 입력 크기, 네트워크 깊이, 데이터셋 특성에 걸쳐 일관된 성능을 보이는 모델을 성공적으로 탐색하였다.
• CAM 시각화 결과, CovidNet3D는 코로나19의 핵심 영상 생물학적 지표인 지상유리 농탁 영역에 집중하고 있음을 확인하여 임상적 해석 가능성을 향상시켰다.
• MosMedData에서의 성능 저하는 NIfTI에서 PNG로의 변환 과정에서의 정보 손실과 클래스 불균형이 원인일 가능성이 높다.
• 훈련 과정은 랜덤 시드에 매우 민감하여, NAS 기반 모델의 강건성은 향후 추가 연구가 필요하다.