mdsajjadullah/FedPrivNet-ChestXray14-CheXpert

FedPrivNet — Privacy-Aware Hybrid Deep Learning Model for Chest X-Ray Classification

<p align="center"> <img src="training_curves.png" width="800"/> </p>

Overview

FedPrivNet is a custom hybrid deep learning architecture designed for multi-label chest X-ray disease classification across 14 thoracic pathologies. The model is specifically engineered as the backbone for a privacy-preserving Federated Learning pipeline, with three core design principles: Differential Privacy compatibility, Demographic Fairness, and Communication Efficiency.

Developed by Md. Sajjad Ullah, this model represents a novel architectural contribution that unifies ResNet-18 and DenseNet-121 feature extraction with purpose-built privacy-safe normalization and a trainable spatial attention mechanism for built-in explainability.

Model Architecture

FedPrivNet is a dual-branch hybrid network with the following novel components:

Key Design Decision: All BatchNorm layers are replaced with GroupNorm throughout the entire network. This makes FedPrivNet natively compatible with DP-SGD (Opacus), which requires per-sample gradient computation — a property that BatchNorm breaks.

Architecture Diagram

Input (224×224 Chest X-Ray) │ ┌────┴────┐ │ │ ResNet-18 DenseNet-121 (256ch) (1024ch) │ │ └────┬────┘ Fusion Conv (256ch, 28×28) │ DPResBlock-1 + SE (256ch) │ DPResBlock-2 + SE (512ch, 14×14) │ DPResBlock-3 + SE (512ch) │ SpatialAttentionGate (novel XAI module) │ GAP → Dropout │ Linear (512→14) │ Sigmoid │ 14-class Output ---

Performance

Overall Metrics

Per-Class AUC-ROC

Spatial Attention Visualization

<p align="center"> <img src="spatialattentionvisualization.png" width="800"/> </p>

Baseline Comparison

FedPrivNet outperforms both individual backbone models when evaluated under identical federated learning conditions (3 training epochs, same dataset partition, same hyperparameters).

The dual-backbone fusion delivers +8.0% AUC improvement over ResNet-18 and +8.7% over DenseNet-121, directly justifying the hybrid architecture design choice.

<p align="center"> <img src="baseline_comparison.png" width="800"/> </p>

Ablation Study

Each architectural component was systematically removed to measure its individual contribution. All variants were trained for 3 epochs from random initialization on the same federated partition.

Key finding: SE Blocks contribute most to performance (+0.024 AUC drop when removed). The Spatial Attention Gate provides additional diagnostic focus and serves as the built-in XAI module. Both components together produce the strongest results.

<p align="center"> <img src="ablation_study.png" width="800"/> </p>

Training Details

Datasets

Usage

Load Pretrained Model

import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download

weights_path = hf_hub_download(
    repo_id  = "mdsajjadullah/FedPrivNet-ChestXray14-CheXpert",
    filename = "FedPrivNet_best.pth"
)

model = FedPrivNet(num_classes=14, dropout=0.3, pretrained=False)
checkpoint = torch.load(weights_path, map_location='cpu',
                        weights_only=False)
model.load_state_dict(checkpoint['model_state'])
model.eval()
print(f"Model loaded | Val AUC: {checkpoint['val_auc']:.4f}")

Inference

from torchvision import transforms
from PIL import Image

DISEASE_LABELS = [
    'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration',
    'Mass', 'Nodule', 'Pneumonia', 'Pneumothorax',
    'Consolidation', 'Edema', 'Emphysema', 'Fibrosis',
    'Pleural_Thickening', 'Hernia'
]

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225]),
])

image = Image.open("chest_xray.png").convert("RGB")
x     = transform(image).unsqueeze(0)

with torch.no_grad():
    probs, attn_map = model(x, return_attention=True)

for label, prob in zip(DISEASE_LABELS, probs[0]):
    if prob > 0.5:
        print(f"  {label}: {prob:.4f}")

Federated Learning Integration

FedPrivNet is the backbone of a full FL pipeline: FedPrivNet Backbone + Opacus 1.4 (DP-SGD, ε ∈ {1, 3, 5, 10}) + Fairness Regularizer (Demographic Parity + Equal Opportunity) + Top-k Gradient Compression (k ∈ {0.30, 0.40, 0.50}) + Grad-CAM + SHAP Explainability Analysis ---

Citation

@misc{ullah2026fedprivnet,
  author    = {Md.Sajjad Ullah},
  title     = {FedPrivNet: Privacy-Aware Hybrid Deep Learning Model
               for Chest X-Ray Classification},
  year      = {2026},
  publisher = {HuggingFace},
  url       = {https://huggingface.co/mdsajjadullah/FedPrivNet-ChestXray14-CheXpert}
}

Author

Md.Sajjad Ullah Department of Computer Science and Engineering University of Asia Pacific, Bangladesh

License

MIT License