Add model code

model.py

@@ -0,0 +1,58 @@
+"""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