Spaces:
Running
Running
File size: 3,691 Bytes
c1969c3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Model 2: DenseNet-121 with transfer learning (CheXNet approach).
Design decisions:
- DenseNet-121: Dense connections promote feature reuse and reduce parameters vs. ResNet.
Chosen because CheXNet (Rajpurkar et al., 2017) demonstrated radiologist-level performance with this backbone.
- ImageNet pretraining: Provides strong low-level feature initialization even for medical images.
Textures, edges, and structural patterns transfer well across domains.
- Two-phase fine-tuning:
Phase 1 — Freeze backbone, train only classifier head (prevents catastrophic forgetting).
Phase 2 — Unfreeze all layers with reduced LR for end-to-end refinement.
- Dual heads: Same rationale as scratch model — shared features, specialized predictions.
"""
from __future__ import annotations
import torch
import torch.nn as nn
from torchvision import models
class CheXVisionDenseNet(nn.Module):
"""DenseNet-121 fine-tuned for chest X-ray pathology detection.
Architecture:
- DenseNet-121 backbone (pretrained on ImageNet)
- Replace final classifier with dual heads:
- Multi-label: 14 pathology classes (sigmoid)
- Binary: Normal vs Abnormal (sigmoid)
"""
def __init__(
self,
num_classes: int = 14,
pretrained: bool = True,
dropout: float = 0.3,
freeze_backbone: bool = False,
) -> None:
super().__init__()
# Load pretrained DenseNet-121
weights = models.DenseNet121_Weights.IMAGENET1K_V1 if pretrained else None
self.backbone = models.densenet121(weights=weights)
# Get the number of features from the classifier
num_features = self.backbone.classifier.in_features
# Remove the original classifier
self.backbone.classifier = nn.Identity()
# Shared feature layer
self.feature_layer = nn.Sequential(
nn.Linear(num_features, 512),
nn.ReLU(inplace=True),
nn.Dropout(p=dropout),
)
# Multi-label classification head (14 pathologies)
self.multilabel_head = nn.Linear(512, num_classes)
# Binary classification head (Normal vs Abnormal)
self.binary_head = nn.Linear(512, 1)
# Optionally freeze backbone (Phase 1 of fine-tuning)
if freeze_backbone:
self.freeze_backbone()
# Initialize new layers
self._init_new_layers()
def _init_new_layers(self) -> None:
for module in [self.feature_layer, self.multilabel_head, self.binary_head]:
for m in module.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
nn.init.constant_(m.bias, 0)
def freeze_backbone(self) -> None:
"""Freeze all backbone parameters (Phase 1: train only classifier heads)."""
for param in self.backbone.parameters():
param.requires_grad = False
def unfreeze_backbone(self) -> None:
"""Unfreeze backbone parameters (Phase 2: end-to-end fine-tuning)."""
for param in self.backbone.parameters():
param.requires_grad = True
def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
"""Forward pass returning both task predictions.
Args:
x: Input tensor of shape (B, 3, 224, 224).
Returns:
Dict with 'multilabel_logits' (B, 14) and 'binary_logits' (B, 1).
"""
features = self.backbone(x)
features = self.feature_layer(features)
return {
"multilabel_logits": self.multilabel_head(features),
"binary_logits": self.binary_head(features),
}
|