File size: 2,083 Bytes
ce593f2 | 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 | """DenseNet121 with an MC-Dropout head for chest X-ray classification.
The backbone is the CheXNet-standard DenseNet121. We replace the classifier
with a Dropout -> Linear head. At inference we keep the Dropout layers active
(see `enable_mc_dropout`) so repeated forward passes give a distribution over
predictions -- the basis for our uncertainty estimate.
"""
import torch
import torch.nn as nn
from torchvision import models
def build_model(num_classes: int = 1, dropout_p: float = 0.3,
pretrained: bool = True) -> nn.Module:
"""DenseNet121 with an MC-Dropout classifier head.
num_classes=1 -> binary task, use with BCEWithLogitsLoss (sigmoid output).
num_classes=N (N>1, multi-label e.g. ChestX-ray14) -> also BCEWithLogitsLoss.
"""
weights = models.DenseNet121_Weights.IMAGENET1K_V1 if pretrained else None
net = models.densenet121(weights=weights)
in_features = net.classifier.in_features
net.classifier = nn.Sequential(
nn.Dropout(p=dropout_p),
nn.Linear(in_features, num_classes),
)
return net
def enable_mc_dropout(model: nn.Module) -> None:
"""Put the model in eval mode but re-activate Dropout layers.
This is what makes MC Dropout work: BatchNorm etc. stay in eval mode
(using running stats), while Dropout keeps sampling.
"""
model.eval()
for m in model.modules():
if isinstance(m, nn.Dropout):
m.train()
@torch.no_grad()
def mc_predict(model: nn.Module, x: torch.Tensor, n_passes: int = 30):
"""Run n stochastic forward passes.
Returns:
mean_prob: (B, C) predictive probability (mean over passes)
uncertainty:(B, C) predictive std over passes (epistemic signal)
"""
enable_mc_dropout(model)
probs = []
for _ in range(n_passes):
logits = model(x)
probs.append(torch.sigmoid(logits))
probs = torch.stack(probs, dim=0) # (T, B, C)
mean_prob = probs.mean(dim=0) # (B, C)
uncertainty = probs.std(dim=0) # (B, C)
return mean_prob, uncertainty
|