Multi-label thoracic pathology classification on the NIH ChestX-ray14 dataset, with Grad-CAM
explanations and a containerised inference API.
Research and educational use only. PulmoNet AI is not a medical device, has not been
clinically validated, and must not be used for diagnosis or treatment decisions.
What this is
An end-to-end pipeline over all 112,120 chest X-rays from 30,805 patients:
Ingest — downloads the dataset, recovers and verifies the official identifiers, and
emits a manifest with patient-level split assignments.
Train — EfficientNet-B4 fine-tuned in PyTorch with class-weighted BCE and dynamic
augmentation, sized to fit a 4 GB laptop GPU.
Explain — Grad-CAM implemented directly against autograd (no wrapper library).
Serve — FastAPI + Docker on AWS Lambda (torch, both endpoints), with a free static
Hugging Face Space as the demo frontend.
The demo is a free static Space; the API is an AWS Lambda container serving both
classification and Grad-CAM. First request after idle takes ~15 s (cold start), then ~200 ms.
pulmonet/ config, data, model, losses, gradcam, engine, inference
scripts/ ingest, train, evaluate, export_onnx
app/ FastAPI app, Lambda handler, local demo UI
space/ static frontend deployed to Hugging Face Spaces
notebooks/ kaggle_train.ipynb - same pipeline on a free cloud GPU
docker/ Dockerfile (generic) and Dockerfile.lambda (AWS)
tests/ 33 tests: data contract, Grad-CAM, API, and backend equivalence
reports/ metrics.json and auroc_table.md (generated, never hand-edited)
Results
See reports/metrics.json and
reports/auroc_table.md. Those files are written by
scripts/evaluate.py and are the only source of truth for numbers quoted about this project.
Evaluation is on the official NIH test split — 25,596 images from 2,797 patients, with no
patient appearing in both training and test.
Model: nih-densenet121 at 224px.
Pathology
AUROC
Wang et al. 2017
Delta
Atelectasis
0.662
0.716
-0.054
Cardiomegaly
0.807
0.807
+0.000
Effusion
0.746
0.784
-0.038
Infiltration
0.632
0.609
+0.023
Mass
0.709
0.706
+0.003
Nodule
0.626
0.671
-0.045
Pneumonia
0.620
0.633
-0.013
Pneumothorax
0.600
0.806
-0.206
Consolidation
0.663
0.708
-0.045
Edema
0.735
0.835
-0.100
Emphysema
0.514
0.815
-0.301
Fibrosis
0.671
0.769
-0.098
Pleural_Thickening
0.681
0.708
-0.027
Hernia
0.735
0.767
-0.032
Mean
0.671
0.738
-0.067
Quickstart
bash
1pip install -r requirements-train.txt
2# install the torch build matching your CUDA runtime, e.g.3pip install torch --index-url https://download.pytorch.org/whl/cu118
45python scripts/ingest.py --limit 200# dry run: verifies the pipeline in ~30 s6python scripts/ingest.py # full ingest, ~15 min, ~1.7 GB on disk78python scripts/train.py --smoke # 50 steps, verifies the training loop9python scripts/train.py # full run10python scripts/evaluate.py # writes reports/1112uvicorn app.main:app --port 7860# http://localhost:786013pytest
Set PULMONET_DATA to control where images and checkpoints go (default C:\pulmonet_data).
Keep it outside any cloud-synced folder.
Training on a free cloud GPU
notebooks/kaggle_train.ipynb runs the same pipeline on
Kaggle's free P100 (16 GB, 30 GPU-hours/week). Two reasons it is the better option than a
4 GB laptop GPU here:
ChestX-ray14 is already hosted on Kaggle, so --source kaggle reads it in place — no
download at all.
16 GB fits batch 32 without gradient accumulation, and nothing competes for VRAM.
Roughly 3–4× faster end to end. scripts/ingest.py --source kaggle is the only piece that
differs; training, evaluation, and reporting are byte-identical.
These are the choices worth understanding; each one is a place where the obvious approach is
wrong.
The split is patient-level, not image-level
A patient contributes multiple X-rays. Splitting randomly by image puts near-duplicate views
of the same chest on both sides of the train/test boundary, which inflates AUROC by several
points and measures memorisation rather than generalisation.
This project uses the official train_val_list.txt / test_list.txt, which partition by
patient — verified empirically: 28,008 training patients and 2,797 test patients, zero
overlap. The validation set is then carved out of the training patients, again by patient ID
(pulmonet/data.py:split_manifest), and tests/test_data.py asserts all three splits are
mutually disjoint.
It is multi-label, not multi-class
One X-ray can show Effusion and Infiltration and Atelectasis simultaneously. So the model
emits 14 independent sigmoids, not a softmax over 14 classes, and trains with
BCEWithLogitsLoss. Probabilities do not sum to 1, and they are not supposed to.
"No Finding" is deliberately not a 15th class — it is the absence of all 14 labels, so an
all-zero prediction vector already expresses it. Modelling it separately would double-count.
The class weights are capped
ChestX-ray14 prevalence ranges from Infiltration at ~17.7% to Hernia at ~0.2%. Unweighted BCE
scores 99.8% "accuracy" on Hernia by always predicting negative, which is worthless — hence
pos_weight = n_neg / n_pos per class.
But Hernia's raw ratio is ~500×, and feeding that in lets a single positive sample dominate the
batch gradient; the loss diverges within a few hundred steps. Clipping at 10× keeps rare classes
meaningfully up-weighted while bounding any one class's contribution
(pulmonet/losses.py:compute_pos_weight).
The augmentation respects anatomy
Horizontal flip, ±10° rotation, mild brightness/contrast, and RandomResizedCrop at 0.85–1.0.
Deliberately no vertical flip and no aggressive rotation: thoracic anatomy has a fixed
orientation and the heart sits on one side, so an upside-down chest X-ray is not a sample the
model will ever see at inference. Augmenting toward it spends capacity on an impossible input.
AUROC, not accuracy
With Hernia at 0.2% prevalence, accuracy is uninformative. AUROC is threshold-free and
prevalence-insensitive, and it is the metric the ChestX-ray14 literature reports, so results
here are directly comparable to Wang et al. (2017), which evaluate.py prints alongside.
Classes with no positive samples in a split report NaN rather than a fabricated 0.5.
Fitting EfficientNet-B4 into 4 GB of VRAM
B4's native resolution is 380px, which OOMs on an RTX 3050 Laptop at any useful batch size. The
working configuration:
Lever
Setting
Effect
Resolution
224
fits; source images are stored at 300px so crops have room
Batch
8
see below — 16 fits only on an idle GPU
Gradient accumulation
×4
effective batch 32 without the memory
AMP (torch.amp)
on
~40% memory saved, ~1.7× faster
channels_last
on
better conv kernel selection
Batch 16 measures at 2.3 GB reserved and looks fine on an idle card — but a running browser
holds ~2 GB of VRAM for compositing, leaving under 2 GB. Training then dies partway through
an epoch with CUDA out of memory (or, more confusingly, CUDNN_STATUS_INTERNAL_ERROR, or
CUBLAS_STATUS_NOT_INITIALIZED — cuDNN and cuBLAS both report allocation failures as opaque
init errors). Batch 8 peaks near 1.25 GB and trains reliably on a machine in use.
Checkpoint selection is on validation mean AUROC, not validation loss — the pos-weighted
loss is dominated by rare classes and can improve while ranking quality is flat.
Grad-CAM is written directly against autograd
pulmonet/gradcam.py is ~60 lines using a forward hook plus a tensor gradient hook on
conv_head. No pytorch-grad-cam dependency, so every line is explicable.
Two non-obvious details it gets right:
Hooks are removed in __exit__. A leaked hook accumulates on every API request.
It forces torch.enable_grad() internally. The API wraps inference in no_grad, under which
a naive implementation silently returns an all-zero map that still looks like a heatmap.
What Grad-CAM is for here: confirming the network attends to lung fields rather than to text
burn-ins, scanner artifacts, or image borders — a documented shortcut-learning failure mode in
chest X-ray models. It is not clinical localisation and not evidence of diagnostic validity.
Ingest verifies its own assumptions
The 300px source mirror ships (image, labels) but drops the original filenames, so patient IDs
and official split membership are not directly available. Its rows are in official split-list
order and its counts match exactly (86,524 / 25,596), so position i of the train split should
be train_val_list.txt[i].
That is an assumption, and a wrong one would silently corrupt every metric in the project. So
scripts/ingest.py re-derives each row's expected findings from Data_Entry_2017.csv and
asserts they match the labels the parquet actually carries — for all 112,120 rows. Any mismatch
aborts the ingest. --source nih-box fetches the original 1024px archives instead, streaming
one at a time so peak disk stays ~13 GB rather than ~90 GB.
Known limitations
Labels are NLP-mined from radiology reports, not expert-annotated. The NIH authors
estimate >90% accuracy, so a few percent of training labels are wrong and the test set
inherits the same noise. This caps achievable AUROC regardless of architecture.
No external validation. Performance on CheXpert, MIMIC-CXR, or any clinical population is
unmeasured; CXR models are known to degrade badly across scanners and sites.
Frontal views only, and no patient metadata (age, sex, prior studies) is used.
Not calibrated. The sigmoid outputs are ranking scores; the pos-weighting deliberately
distorts them away from true posterior probabilities. Do not read them as "X% chance of
disease."
Dataset and citation
Wang X., Peng Y., Lu L., Lu Z., Bagheri M., Summers R.M. ChestX-ray8: Hospital-scale Chest
X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common
Thorax Diseases. CVPR 2017.
