Deep X-Ray Vision
흉부 X-ray 멀티라벨 분류에서 불균형 대응(ASL) · 학습 안정화(EMA) · 추론 최적화(TTA) · 모델 앙상블을 조합해 성능을 끌어올린 대회형 프로젝트입니다.
소개
본 프로젝트는 흉부 X-ray 멀티라벨 분류 대회에서 단순히 모델을 구현하는 것을 넘어,성능이 정체되는 구간을 어떻게 돌파할 것인가에 집중한 실험 중심 프로젝트입니다. 다양한 모델 구조를 비교하며 성능 변화를 분석했고, 단일 모델의 한계를 확인한 뒤 불균형 완화(ASL), 학습 안정화(EMA), 추론 보정(TTA), 그리고 앙상블 전략을 단계적으로 적용해 최종 성능을 개선했습니다.
모델 탐색 과정
EfficientNet(B3/B4/B5), ResNet(50/101), ViT(384/512) 등 다양한 아키텍처를 실험하며 성능을 비교했습니다. 또한 HuggingFace pretrained 모델을 활용해 전이학습 기반 성능 향상도 시도했습니다. 그러나 단일 모델만으로는 점수 상승이 제한되는 구간에 도달했고, 서로 다른 특성을 가진 CNN과 Transformer 모델을 결합한 앙상블 전략이 가장 안정적인 성능 향상을 만든다는 결론에 도달했습니다.
핵심 기능
불균형 완화 (ASL)
Multi-label 데이터에서 음성 샘플 편향을 줄이고 희소 클래스 학습을 강화하도록 손실 함수를 개선했습니다.
학습 안정화 (EMA)
가중치 이동 평균을 적용해 validation 변동을 줄이고 일반화 성능을 안정화했습니다.
추론 보정 (TTA)
좌우 반전 기반 Test-Time Augmentation으로 예측을 평균화해 추론 안정성을 높였습니다.
모델 앙상블
EfficientNet-B4와 ViT-384를 결합해 서로 다른 모델 특성을 보완했습니다.
모델 예측 결과
Input X-ray
Prediction + Grad-CAM
최종 학습된 모델에 X-ray 이미지를 입력하여, 가장 높은 확률로 예측된 질환 결과를 시각화한 예시입니다.
문제점 & 해결방법
문제 — Multi-label 데이터 불균형: 'No Finding' 비중이 과도해 모델이 음성(정상) 예측에 치우치고, 희소 클래스의 재현율이 잘 오르지 않았습니다.
해결 — ASL(Asymmetric Loss)을 도입해 쉬운 음성 샘플의 손실 기여를 줄이고, 어려운/희소 클래스 샘플에 더 집중하도록 학습을 유도했습니다. 또한 클래스별 AUC를 함께 추적해, 전체 점수 상승이 특정 클래스만의 개선이 아닌지 지속적으로 확인했습니다.
문제 — 입력 사이즈 불일치: EfficientNet-B4(380)와 ViT-384(384)의 기본 입력 크기가 달라, 학습/추론 파이프라인이 꼬이거나 앙상블 결과가 흔들릴 수 있었습니다.
해결 — 전처리 규칙을 ‘기준 해상도 384’로 통일하고, 모델별로 필요한 사이즈(380/384)는 추론 단계에서만 안전하게 리사이즈하도록 분리했습니다. 이렇게 해서 데이터 로딩/증강은 한 번만 수행하고, 모델 입력만 유연하게 맞춰 앙상블 일관성을 확보했습니다.
문제 — PyTorch(torch_xla)로 TPU를 적용할 때 GPU와 다른 실행 구조 때문에 학습/추론 안정성 문제와 저장 오류가 발생했습니다.
해결 — TPU 실행 방식에 맞게 학습/추론 파이프라인을 조정하고, 결과 저장 오류를 해결해 TPU에서도 안정적으로 추론이 가능하도록 개선했습니다.
결과/성과
ASL로 불균형 완화 · EMA로 학습 안정화 · TTA와 앙상블로 최종 성능 향상