Chest X-ray Pneumonia Classifier with MC-Dropout Uncertainty
A DenseNet121 chest X-ray classifier (CheXNet-style backbone) that reports a per-image uncertainty score alongside its prediction, and uses that uncertainty to drive a selective-referral (risk-coverage) analysis: when the model defers its most-uncertain cases to a radiologist, AUROC on the cases it keeps goes up.
Model description
- Backbone: DenseNet121, ImageNet-pretrained.
- Head:
Dropout(p=0.3) -> Linear(in_features, 1), kept active at inference time (Monte Carlo Dropout) so repeated stochastic forward passes give a distribution over predictions. - Task: binary classification, chest X-ray -> {NORMAL, PNEUMONIA}.
- Training data: Kaggle Chest X-Ray Images (Pneumonia) dataset (Kermany et al.), 5,219 train / 19 val / 625 test images.
- Training: 5 epochs, AdamW (lr=1e-4, weight_decay=1e-4),
BCEWithLogitsLoss, on a single NVIDIA B200 GPU (HiPerGator).
Output
The model returns a single logit. torch.sigmoid(logit) gives the
probability that the image is class index 1 (PNEUMONIA). Running
mc_predict() (30 stochastic forward passes with dropout active) returns two
values per image instead of one:
- mean probability across passes (the prediction)
- std across passes (the epistemic uncertainty score)
This model does not output segmentation masks, bounding boxes, or localization — whole-image binary classification only.
Results
- Full test set: AUROC 0.967, ECE 0.320
- Uncertainty tracks error: mean MC-dropout uncertainty is 0.008 on correct predictions vs. 0.030 on wrong ones (~3.7x higher).
- Selective referral: AUROC on the retained set rises from 0.967 at full
coverage to 0.987 when the most-uncertain 30% of cases are referred away.
See
referral_curve.png.
Honest limitations
- ECE of 0.320 indicates the model is overconfident. This doesn't undermine the referral result (which relies on relative uncertainty ranking, not calibrated probabilities), but it matters for any real deployment claim — temperature scaling would be a natural next step.
- The validation split used during training is tiny (19 images) and saturates at AUROC 1.000, so it's a weak model-selection signal; the 625-image test set is the trustworthy number.
- Pneumonia detection (binary) is an easier task than multi-label nodule/mass detection (e.g. NIH ChestX-ray14); this is meant as a methodology demonstration, not a diagnostic-grade nodule detector.
- MC Dropout captures epistemic uncertainty only; deep ensembles are a stronger (heavier) baseline.
- No PHI is involved (public de-identified dataset), but real clinical use would require DICOM handling, de-identification, and IRB review.
How to use
import torch
from model import build_model, mc_predict # from this repo
from data import _tf # preprocessing transform
model = build_model(num_classes=1)
model.load_state_dict(torch.load("ckpt.pt", map_location="cpu"))
# x: a batch of preprocessed images, shape (B, 3, 224, 224)
mean_prob, uncertainty = mc_predict(model, x, n_passes=30)
Files
ckpt.pt— trained model weights (state dict).model.py— DenseNet121 + MC-dropout head;mc_predictfor sampled inference.data.py— preprocessing transforms (grayscale->3ch, resize, ImageNet norm).referral_curve.png— AUROC-vs-coverage plot (the headline result).