ChestX6-SSL-Benchmark Checkpoints

This repository hosts the pretrained and fine-tuned model checkpoints accompanying the manuscript:

"A Deployment Risk Score Framework for Self-Supervised Chest X-Ray Classification: Calibrated Multi-Objective Evaluation Under Annotation Scarcity and Distribution Shift"

Authors: Kashif Mahmood, Romana Aziz, Muhammad Ramzan, Mahwish Ilyas, and Ala Saleh Alluhaidan

Research use only. These models are not medical devices, have not been validated for clinical use, and must not be used to inform diagnosis or treatment.

Overview

The checkpoints correspond to the experiments performed on the ChestX6 dataset and include:

Five supervised baseline models (including ViT-S/16 trained from scratch, the architecture-matched control for MAE)
SimCLR self-supervised pretrained encoder
MAE-style pretrained encoder (patch-embedding layer only; see the note below)
SimCLR and MAE-style encoders pretrained on the training partition only (ssl_trainonly_data/), used for the evaluation-independence sensitivity analysis
Fine-tuned SimCLR models
Fine-tuned MAE models

Models are provided for 10%, 20%, and 100% labeled-data settings with three random seeds (42, 123, 456) for each configuration. SimCLR and MAE were each pretrained once; only fine-tuning was repeated across seeds.

Model	Pretraining	Folder
ResNet50 (Scratch)	None	`supervised/`
ResNet50 (ImageNet)	Supervised ImageNet	`supervised/`
EfficientNet-B0	Supervised ImageNet	`supervised/`
MobileViT-XS	Supervised ImageNet	`supervised/`
ViT-S/16 (Scratch)	None (architecture-matched control for MAE)	`supervised/`
SimCLR + ResNet50	Contrastive SSL (ChestX6)	`finetuned/` (encoder in `ssl/`; training-partition-only encoder in `ssl_trainonly_data/`)
MAE + ViT-S/16	MAE-style SSL (ChestX6); only the patch-embedding layer was pretrained	`finetuned/` (encoder in `ssl/`; training-partition-only encoder in `ssl_trainonly_data/`)

What the "MAE" checkpoints contain. In the pretraining code used for the paper the visible patch embeddings were passed directly to the decoder, so the ViT's transformer blocks, class token, positional embedding and final norm were not part of the pretraining forward pass and were not updated. The mae_vits16 checkpoints therefore have a pretrained patch-embedding layer and randomly initialised transformer blocks, as in vit_s16_scratch. They do not represent masked-reconstruction pretraining of a full ViT encoder. A corrected implementation (not used for the paper's results) is in the GitHub repository under scripts/. For SimCLR, the NT-Xent loss was computed inside micro-batches of 64 images (126 negatives per anchor), not over the effective batch of 256.

Repository Structure

supervised/
    ResNet50 (Scratch)
    ResNet50 (ImageNet)
    EfficientNet-B0
    MobileViT-XS
    ViT-S/16 (Scratch)

ssl/
    SimCLR pretrained encoder (.pth)
    MAE-style pretrained encoder (.pth; patch-embedding layer pretrained, see note above)
    SimCLR ResNet50 Encoder (.pth)
    SimCLR log loss (.json)

ssl_trainonly_data/
    MAE-style pretrained encoder - training data only (.pth; patch-embedding layer pretrained, see note above)
    SimCLR pretrained encoder - training data only (.pth)
    SimCLR log loss (.json)

finetuned/
    SimCLR fine-tuned models (.pt)
    MAE fine-tuned models (.pt)

Supervised and fine-tuned checkpoints are named {model_name}_{subset}_seed{seed}_best.pt, where subset identifies the label budget (10%, 20% or 100%) and model_name is one of resnet50_scratch, resnet50_pretrained, efficientnet_b0, mobilevit_xs, vit_s16_scratch, simclr_resnet50 or mae_vits16.

Dataset

The experiments use ChestX6, a six-class chest X-ray benchmark (Covid-19, Emphysema, Normal, Pneumonia-Bacterial, Pneumonia-Viral, Tuberculosis). ChestX6 is a compilation of two independently verified public sources; it is not hosted as a standalone dataset. No images are included in this repository.

Covid-19, Normal, Pneumonia-Bacterial, Pneumonia-Viral, Emphysema (82.1%): https://www.kaggle.com/datasets/minhnhat232/dataset-covid-bacterial-viral-normal-emphysema
Tuberculosis (17.6%): https://www.kaggle.com/datasets/tawsifurrahman/tuberculosis-tb-chest-xray-dataset

A further 0.3% matched the curated dataset of Sait et al. and the COVID-19 Radiography Database. The source datasets remain under their providers' terms; see the licence table in the GitHub repository (data/README.md).

18,036 images were collected. During preprocessing, 48 exact duplicates were removed (MD5 audit), resulting in 17,988 unique images used in all experiments (split 70/10/20, seed 42, stratified).

Cross-dataset evaluation uses ChestMNIST (28×28 release, binary labels, derived from NIH ChestX-ray14) through the MedMNIST package. It is used as an external test set only. NIH ChestX-ray14 was provided by the NIH Clinical Center.

Loading

The checkpoints store PyTorch model weights (torch.load(path)["model_state"] for supervised and fine-tuned models). Use the pinned revision below so you get exactly the files used in the paper:

import torch, timm, torchvision
from huggingface_hub import list_repo_files, hf_hub_download

REPO = "Kashif-Mahmood007/chestx6-ssl-checkpoints"
REV  = "1764b2409e06a91e81eda0f8bfe0c8ae4e199a38"          # the revision cited in the paper

print(list_repo_files(REPO, revision=REV))   # every file at this revision

# Pick a file from the listing (adjust the filter to choose the model, budget and seed)
files = list_repo_files(REPO, revision=REV)
fname = next(f for f in files if "simclr_resnet50" in f and "seed42" in f and f.endswith(".pt"))
path  = hf_hub_download(REPO, filename=fname, revision=REV)


def build(model_name, num_classes=6):
    """Architectures as defined in the repository notebooks (NB2/NB3)."""
    if model_name in ("resnet50_scratch", "resnet50_pretrained"):
        m = torchvision.models.resnet50(weights=None)
        m.fc = torch.nn.Linear(2048, num_classes)
        return m
    if model_name == "efficientnet_b0":
        return timm.create_model("efficientnet_b0", pretrained=False, num_classes=num_classes)
    if model_name == "mobilevit_xs":
        return timm.create_model("mobilevit_xs", pretrained=False, num_classes=num_classes)
    if model_name == "vit_s16_scratch":
        return timm.create_model("vit_small_patch16_224", pretrained=False, num_classes=num_classes)
    if model_name == "simclr_resnet50":
        enc = torchvision.models.resnet50(weights=None)
        enc.fc = torch.nn.Identity()
        return torch.nn.Sequential(enc, torch.nn.Linear(2048, num_classes))
    if model_name == "mae_vits16":
        enc = timm.create_model("vit_small_patch16_224", pretrained=False, num_classes=0)
        return torch.nn.Sequential(enc, torch.nn.Linear(384, num_classes))
    raise ValueError(model_name)


model = build("simclr_resnet50")
model.load_state_dict(torch.load(path, map_location="cpu")["model_state"])
model.eval()

Refer to the accompanying GitHub repository for the complete training pipeline, preprocessing and normalisation, configuration files, and evaluation code.

Intended use and limitations

Intended for research on label efficiency, robustness, calibration and model selection. Not a medical device and not validated for clinical use.
Retrospective comparison on public datasets; results do not establish performance on any hospital's data.
SimCLR and MAE pretraining each ran once, so pretraining-to-pretraining variance is not captured; seed variation reflects fine-tuning only.
The SimCLR and MAE pipelines differ in backbone, objective, epochs and updates and are not compute-matched; differences between them are differences between pipelines, not effects of the pretraining objective. The MAE pipeline pretrained only the patch-embedding layer (see the note above).
Robustness was tested against eight synthetic image perturbations, and external-dataset transfer against ChestMNIST (28×28 images, upsampled; binary labels). Neither isolates the effect of scanner or acquisition differences.
Checkpoints were trained on public datasets whose licences differ from this repository's; check that your intended use is compatible with the source terms.

GitHub Repository

The complete source code, training pipeline, notebooks, configuration files, dataset split definitions, and reproduction instructions are available on GitHub:

Repository: ChestX6-SSL-Benchmark

The archived release used for the paper is available at Zenodo: https://doi.org/10.5281/zenodo.23288310 (release v1.1.0-scirep).

The GitHub repository includes:

End-to-end training and evaluation notebooks
Fixed train/validation/test split definitions, MD5 checksums and the group-respecting Tuberculosis split
Experiment results and manuscript figures
Instructions for reproducing all experiments reported in the accompanying manuscript

Persistent identifiers

Checkpoints (this repository), revision used in the paper: 1764b2409e06a91e81eda0f8bfe0c8ae4e199a38
Checkpoints DOI: https://doi.org/10.57967/hf/10799
Code and results (release v1.1.0-scirep): https://doi.org/10.5281/zenodo.23288310
Split indices and checksums (Kaggle, version 4): https://doi.org/10.34740/KAGGLE/DSV/20558883

Citation

If you use these checkpoints, please cite the associated manuscript.

@unpublished{mahmood2026drs,
  title  = {A Deployment Risk Score Framework for Self-Supervised Chest X-Ray Classification: Calibrated Multi-Objective Evaluation Under Annotation Scarcity and Distribution Shift},
  author = {Kashif Mahmood and Romana Aziz and Muhammad Ramzan and Mahwish Ilyas and Ala Saleh Alluhaidan},
  note   = {Manuscript under revision at Scientific Reports},
  year   = {2026}
}

License

MIT License (code and model checkpoints). No image data is included. The datasets used for training remain under their providers' terms, and some of those terms may restrict derivative or commercial use of models trained on the data: check them before any use beyond research on this benchmark.

Downloads last month: -; Downloads are not tracked for this model. How to track