[논문 리뷰] GANs 'N Lungs: improving pneumonia prediction

Tatiana Malygina, Elena Ericheva외 1명|arXiv (Cornell University)|2019. 08. 01.

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

한 줄 요약

이 논문은 고도로 불균형한 X-ray 데이터셋에서 폐렴 예측 성능을 향상시키기 위해 CycleGAN 기반의 데이터 증강 방법을 제안한다. 정상과 폐렴 클래스 간의 이미지 전환을 통해 합성 이미지를 생성함으로써 훈련 데이터를 균형 잡고 분류기 성능을 향상시킨다. 검증 데이터에서 AUC ROC는 0.9929, AUC PR는 0.9865를 기록한다.

ABSTRACT

We propose a novel method to improve deep learning model performance on highly-imbalanced tasks. The proposed method is based on CycleGAN to achieve balanced dataset. We show that data augmentation with GAN helps to improve accuracy of pneumonia binary classification task even if the generative network was trained on the same training dataset.

연구 동기 및 목표

• 정상 케이스가 폐렴 케이스에 비해 크게 많아지는 고도로 불균형한 의료 영상 데이터셋 문제를 해결한다.
• 훈련 데이터를 균형 잡기 위해 현실적이며 클래스 간 전환된 X-ray 영상을 생성하는 데이터 증강 기법을 개발한다.
• 광범위한 계산 자원이 필요 없이도 딥 러닝 모델의 이진 폐렴 분류 성능을 향상시킨다.
• 제한된 하드웨어 자원을 가진 임상 기계 학습 파이프라인에 실용적으로 데이터 증강을 적용할 수 있도록 한다.

제안 방법

• 정상 X-ray를 폐렴 유사 이미지로, 폐렴 X-ray를 정상 유사 이미지로 변환하는 두 개의 생성기 G_{0→1}과 G_{1→0}을 학습하기 위해 CycleGAN과 사이클 일致 손실을 사용한다.
• 원본 데이터셋에서 생성기를 훈련시켜 도메인 간 번역을 학습하면서 해부학적 구조와 병리학적 특징을 유지한다.
• 원본 이미지와 생성된 이미지를 조합하여 각 클래스에 동일한 샘플 수를 가지도록 균형 잡힌 훈련 세트를 구성한다.
• 이중 교차 엔트로피 손실을 사용하여 증강된 데이터셋에서 DenseNet-121 분류기를 미세 조정한다. 폐렴 대 정상 분류를 위한 것이다.
• 주의 영역을 시각화하고 모델의 해석 가능성 여부를 검증하기 위해 클래스 활성화 맵을 적용한다.
• 모델 간 공정한 비교를 위해 CheXNet에서 사용된 동일한 검증 분할을 사용한다.

실험 결과

연구 질문

• RQ1 CycleGAN 기반의 데이터 증강이 고도로 불균형한 X-ray 데이터셋을 폐렴 분류에 효과적으로 균형 잡을 수 있는가?
• RQ2 정상과 폐렴 클래스 간 전환된 합성 이미지를 생성함으로써 분류기 성능이 향상되는가?
• RQ3 원본 데이터셋 대비 증강된 데이터셋에서 훈련된 폐렴 분류기 성능은 어떻게 변화하는가?
• RQ4 단일 GPU와 같은 제한된 계산 자원에서도 이 방법이 효율적으로 적용될 수 있는가?
• RQ5 생성된 이미지는 임상적 타당성과 모델의 강건성 측면에서 실제 이미지와 어떻게 비교되는가?

주요 결과

• 제안된 방법은 검증 세트에서 AUC ROC를 기준선의 0.9745에서 0.9929로 상승시켜 뚜렷한 성능 향상을 보였다.
• AUC PR는 0.9580에서 0.9865로 향상되어 소수의 폐렴 클래스에서 더 강력한 성능을 나타냈다.
• 추가로 사전 훈련된 데이터셋을 사용함으로써 AUC ROC는 0.9939로, AUC PR는 0.9883으로 더욱 향상되었다.
• 클래스 활성화 맵을 통해 증강된 모델이 임상적으로 관련 있는 영역에 주목하는 것을 확인하여 해석 가능성 향상을 입증했다.
• 단일 NVIDIA 1070 GPU로도 높은 성능를 달성하여 계산 오버헤드가 낮음을 시사했다.
• 생성된 이미지는 시각적으로 타당성이 있으며, 서로 구별 가능했으며, 임상 평가 가능성까지 제기되었다.