[논문 리뷰] Study and development of a Computer-Aided Diagnosis system for classification of chest x-ray images using convolutional neural networks pre-trained for ImageNet and data augmentation

Vinicius Pavanelli Vianna|arXiv (Cornell University)|2018. 06. 03.

COVID-19 diagnosis using AI참고 문헌 15인용 수 13

한 줄 요약

이 연구는 ImageNet에서 미리 훈련된 ConvNet(AlexNet, ResNet, Inception, SqueezeNet)과 데이터 증강을 사용하여 흉부 X선 영상의 정상 및 폐렴 카테고리로 분류하는 컴퓨터 지원 진단(CAD) 시스템을 개발한다. 주요 발견은 전이 학습과 데이터 증강을 조합함으로써 과적합을 크게 감소시키고, ResNet18을 사용하여 최고의 검증 정확도 96.37%를 달성했다.

ABSTRACT

Convolutional neural networks (ConvNets) are the actual standard for image recognizement and classification. On the present work we develop a Computer Aided-Diagnosis (CAD) system using ConvNets to classify a x-rays chest images dataset in two groups: Normal and Pneumonia. The study uses ConvNets models available on the PyTorch platform: AlexNet, SqueezeNet, ResNet and Inception. We initially use three training styles: complete from scratch using random initialization, using a pre-trained ImageNet model training only the last layer adapted to our problem (transfer learning) and a pre-trained model modified training all the classifying layers of the model (fine tuning). The last strategy of training used is with data augmentation techniques that avoid over fitting problems on ConvNets yielding the better results on this study

연구 동기 및 목표

• 딥 러닝을 활용하여 신뢰성 있고 계산적으로 효율적인 흉부 X선 영상 분류를 위한 CAD 시스템을 개발하기.
• 작은 비균형 데이터셋에서 전이 학습과 데이터 증강이 분류 정확도 향상에 얼마나 효과적인지 평가하기.
• 랜덤 초기화, 전이 학습, 미세조정, 데이터 증강 등의 다양한 훈련 전략을 ConvNet 모델에 적용하여 비교하기.
• 제한된 의료 영상 데이터로 훈련할 때 딥 러닝 모델의 과적합을 줄이기.

제안 방법

• PyTorch 플랫폼에서 제공하는 ImageNet 가중치로 초기화된 사전 훈련된 ConvNet(AlexNet, ResNet18, Inception, SqueezeNet)을 사용하였다.
• 세 가지 훈련 전략을 구현: 훈련에서부터 시작하기, 전이 학습(최종 레이어만 미세조정하기), 분류기 전체에 대한 전면적인 미세조정.
• PyTorch의 torchvision.transforms를 활용해 무작위 크기 조정, 자르기, 수평 뒤집기 등을 통해 데이터 증강을 구현하여 훈련 데이터의 다양성을 높였다.
• 모든 실험에서 빠른 수렴과 낮은 손실을 보이는 점을 고려해 기본 초모수를 사용한 Adam 옵timizer를 사용하였다.
• 정상 및 폐렴 클래스 간의 균형을 맞추기 위해 훈련 및 테스트 세트를 1:1 비율로 조정하였다.
• 정확도, 손실 곡선, 혼동 행렬, ROC 곡선을 사용하여 모델 성능을 평가하였다.

실험 결과

연구 질문

• RQ1 작은 흉부 X선 영상 데이터셋에서 ImageNet 사전 훈련 가중치를 사용한 전이 학습이 훈련에서부터 시작하는 것보다 어떻게 다를까?
• RQ2 제한된 의료 영상 데이터로 훈련된 ConvNet에서 데이터 증강이 과적합을 얼마나 줄이는가?
• RQ3 모델 아키텍처와 훈련 전략의 어떤 조합이 흉부 X선 진단에서 가장 높은 분류 정확도를 달성하는가?
• RQ4 혼동 행렬과 ROC 곡선은 데이터 증강을 적용한 경우와 적용하지 않은 경우의 분류기의 강건성을 어떻게 반영하는가?

주요 결과

• 전이 학습과 데이터 증강을 함께 사용한 ResNet18 모델이 가장 높은 검증 정확도 96.37%를 기록했다.
• 동일한 데이터 증강 및 미세조정 전략 하에서 Inception 모델은 검증 정확도 95.51%를 달성했다.
• 훈련에서부터 시작한 경우 심각한 과적합이 발생하여 훈련 정확도는 100%에 도달했지만, 검증 정확도는 86.38%로 정체되었다.
• 데이터 증강은 낮고 안정적인 검증 손실과 향상된 일반화 능력을 통해 과적합을 크게 감소시켰다.
• 혼동 행렬을 통해 데이터 증강 적용 시 거짓 양성과 거짓 음성 수가 크게 감소한 것으로 나타났다.
• ROC 곡선을 통해 데이터 증강이 클래스 간 분리 능력을 향상시켜 AUC를 이상적인 북서쪽 모서리에 더 가깝게 만들었으며, 이는 분류기의 강건성이 향상되었음을 시사한다.