Pneumonia Classification Model (EfficientNet-B0)
Binary classifier for chest X-ray images: Pneumonia vs Normal.
Architecture
- Model: EfficientNet-B0 (timm) — 4.0M parameters
- Input: 224×224 RGB (grayscale CXR → 3-channel)
- Pretraining: ImageNet
- Head: 2-class linear classifier
Dataset
- Source: hf-vision/chest-xray-pneumonia
- Splits: Train (5,216) | Validation (16) | Test (624)
- Class imbalance: ~1:2.9 (Normal:Pneumonia in train)
Training Recipe
- Reproducibility: Fixed seed = 42, deterministic mode
- Augmentation: RandomHorizontalFlip, RandomRotation(15°), ColorJitter(brightness/contrast 0.05)
- Normalization: ImageNet stats
- Loss: Weighted CrossEntropy (inverse class frequency)
- Sampler: WeightedRandomSampler to balance batches
- Optimizer: AdamW (lr=1e-4, weight_decay=1e-4)
- Epochs: 5 (stratified 200+200 subset for balanced training)
Test Performance
| Metric | Score |
|---|
| Accuracy | 0.8125 |
| Precision (Pneumonia) | 0.7910 |
| Recall (Pneumonia) | 0.9513 |
| F1-Score | 0.8638 |
| ROC-AUC | 0.9037 |
Explainability
Grad-CAM heatmaps are included in gradcam/ to visualize regions influencing predictions.
Files
model.pt — Trained model checkpoint (state_dict + config + results)results.json — Detailed metrics and class distributioncm.png — Confusion matrixroc.png — ROC curvegradcam/*.png — Grad-CAM overlays
Limitations
- Trained on a small stratified subset (400 images) due to compute constraints; full-dataset training would further improve generalization.
- Validation set is very small (16 images); results may have high variance.
- No clinical validation performed; intended for research/educational use only.
- Class imbalance in the original dataset was addressed via sampling but may not fully represent real-world prevalence.
