Address CodeRabbit review: validate threshold, restore train state, drop invalid labels

- set_threshold: reject non-scalar qhat (prevents silent broadcasting bugs)
- calibrate: save and restore network.training state (avoid side effect of eval())
- accumulate: drop invalid labels (negative or >= C) instead of clamping them
  (clamping silently remapped -1 -> 0 and 255 -> C-1, corrupting LAC scores)
- accumulate: move scores to CPU after detach (avoid GPU OOM on large 3D cal sets)
- tests: add coverage for scalar validation, train-state restore, invalid-label
  handling; fix assertFalse(.all()) -> .any() for excluded-class assertions
- docstrings: add Returns/Raises sections to _quantile_threshold, accumulate,
  calibrate, set_threshold, __call__

Signed-off-by: Colin Son <txmed82@users.noreply.github.com>

Author: txmed82

monai/inferers/conformal_predictor.py

@@ -32,6 +32,16 @@ def _quantile_threshold(scores: torch.Tensor, alpha: float) -> torch.Tensor:
     (Vovk et al. 2005; Angelopoulos & Bates 2021). ``scores`` is a 1-D tensor of
     non-conformity scores from the held-out calibration split; the returned scalar
     lives on the same device/dtype as ``scores``.
+
+    Args:
+        scores: 1-D tensor of non-conformity scores from the calibration split.
+        alpha: mis-coverage level in ``(0, 1)``.
+
+    Returns:
+        Scalar tensor ``qhat`` on the same device/dtype as ``scores``.
+
+    Raises:
+        ValueError: if ``scores`` is empty or ``alpha`` is not in ``(0, 1)``.
     """
     n = scores.numel()
     if n <= 0:
@@ -90,11 +100,15 @@ def accumulate(self, probs: torch.Tensor, labels: torch.Tensor) -> None:
             probs: softmax probabilities ``(B, C, spatial...)`` in ``[0, 1]`` summing to 1 over C.
             labels: integer class indices ``(B, 1, spatial...)`` or ``(B, spatial...)`` with values
                 in ``[0, C)``. Shape must broadcast against ``probs`` spatial dims.
+
+        Raises:
+            ValueError: if ``probs`` has fewer than 2 dimensions.
         """
         if probs.ndim < 2:
             raise ValueError(f"probs must be (B, C, spatial...), got shape {tuple(probs.shape)}.")
+        c = probs.shape[1]
         # flatten to (N, C)
-        probs_flat = probs.reshape(probs.shape[0], probs.shape[1], -1).movedim(1, -1).reshape(-1, probs.shape[1])
+        probs_flat = probs.reshape(probs.shape[0], c, -1).movedim(1, -1).reshape(-1, c)
         labels_flat = labels.reshape(-1).long()
         if not self.include_background:
             # drop background-labeled voxels (class 0) from calibration so the threshold isn't
@@ -103,12 +117,26 @@ def accumulate(self, probs: torch.Tensor, labels: torch.Tensor) -> None:
             keep = labels_flat != 0
             probs_flat = probs_flat[keep]
             labels_flat = labels_flat[keep]
-        labels_flat = labels_flat.clamp(min=0, max=probs_flat.shape[1] - 1)
+        # reject invalid labels (negative or >= C) outright rather than silently clamping them,
+        # which would corrupt the non-conformity scores.
+        valid = (labels_flat >= 0) & (labels_flat < c)
+        probs_flat = probs_flat[valid]
+        labels_flat = labels_flat[valid]
+        if labels_flat.numel() == 0:
+            return  # nothing to accumulate this batch (all labels invalid or all bg-excluded)
         true_p = probs_flat.gather(1, labels_flat.unsqueeze(1)).squeeze(1)
-        self._scores.append((1.0 - true_p).detach())
+        # move to CPU to avoid GPU OOM on large per-voxel calibration sets
+        self._scores.append((1.0 - true_p).detach().cpu())
 
     def calibrate(self) -> torch.Tensor:
-        """Return the split-conformal threshold ``qhat`` from all accumulated scores."""
+        """Return the split-conformal threshold ``qhat`` from all accumulated scores.
+
+        Returns:
+            Scalar tensor ``qhat`` on CPU.
+
+        Raises:
+            RuntimeError: if no calibration scores have been accumulated.
+        """
         if not self._scores:
             raise RuntimeError("No calibration scores accumulated; call accumulate(probs, labels) first.")
         all_scores = torch.cat(self._scores)
@@ -174,9 +202,19 @@ def __init__(
             self.set_threshold(qhat)
 
     def set_threshold(self, qhat: torch.Tensor) -> None:
-        """Set (or update) the calibrated threshold. Lets you keep one inferer and re-calibrate."""
+        """Set (or update) the calibrated threshold. Lets you keep one inferer and re-calibrate.
+
+        Args:
+            qhat: scalar ``torch.Tensor`` threshold.
+
+        Raises:
+            TypeError: if ``qhat`` is not a ``torch.Tensor``.
+            ValueError: if ``qhat`` is not a scalar (has more than one element).
+        """
         if not isinstance(qhat, torch.Tensor):
             raise TypeError(f"qhat must be a torch.Tensor, got {type(qhat)}.")
+        if qhat.numel() != 1:
+            raise ValueError(f"qhat must be a scalar tensor, got {qhat.numel()} elements.")
         self.qhat = qhat.detach().clone()
 
     def calibrate(
@@ -193,10 +231,14 @@ def calibrate(
 
         Returns:
             The calibrated ``qhat`` (also stored and used by subsequent ``__call__`` invocations).
+
+        Raises:
+            TypeError: if the network returns a non-Tensor.
         """
         if device is None:
             param = next(network.parameters(), None)
             device = param.device if param is not None else torch.device("cpu")
+        was_training = getattr(network, "training", False)
         network.eval()
         cal = ConformalCalibrator(alpha=self.alpha, score=self.score, include_background=self.include_background)
         iterator = cal_loader
@@ -215,6 +257,8 @@ def calibrate(
                 cal.accumulate(probs.to(device), batch["label"].to(device))
         qhat = cal.calibrate()
         self.set_threshold(qhat)
+        if was_training:
+            network.train()
         return qhat
 
     def __call__(
@@ -230,6 +274,10 @@ def __call__(
         Returns:
             sets: bool tensor ``(B, C, spatial...)``, ``True`` where class ``c`` is in the set.
             For the underlying softmax, call ``network(inputs).softmax(1)``.
+
+        Raises:
+            RuntimeError: if no threshold has been set.
+            TypeError: if the network returns a non-Tensor.
         """
         if self.qhat is None:
             raise RuntimeError("No threshold set; call set_threshold(qhat) or calibrate(...) first.")

tests/inferers/test_conformal_predictor.py

@@ -109,6 +109,25 @@ def test_calibrate_empty_raises(self):
         with self.assertRaises(RuntimeError):
             ConformalCalibrator().calibrate()
 
+    def test_invalid_labels_are_dropped_not_clamped(self):
+        # labels -1 and 99 (out of range for C=3) must be dropped, not clamped to 0/2.
+        # If clamped, score for "label 99" would be 1 - softmax[2] = 0.9, corrupting qhat.
+        probs = torch.tensor([[[0.8, 0.8], [0.1, 0.1], [0.1, 0.1]]])  # (1, 3, 2)
+        labels = torch.tensor([[[-1, 99]]])  # both invalid
+        cal = ConformalCalibrator(alpha=0.1)
+        cal.accumulate(probs, labels)
+        with self.assertRaises(RuntimeError):
+            # all labels dropped -> no scores -> calibrate raises
+            cal.calibrate()
+
+    def test_mixed_valid_invalid_labels_keeps_valid(self):
+        # one valid label (class 0, score 0.2) + one invalid (class 99, dropped).
+        # qhat should reflect only the valid score.
+        probs = torch.tensor([[[0.8, 0.8], [0.1, 0.1], [0.1, 0.1]]])  # (1, 3, 2)
+        labels = torch.tensor([[[0, 99]]])  # voxel0 valid, voxel1 invalid
+        qhat = self._cal_batch(probs, labels)
+        assert_allclose(qhat, torch.tensor(0.2), atol=1e-6)
+
 
 class TestConformalPredictor(unittest.TestCase):
     def test_set_and_predict(self):
@@ -128,8 +147,8 @@ def net(x):
         self.assertEqual(sets.shape, (1, 3, 2, 2))
         self.assertEqual(sets.dtype, torch.bool)
         self.assertTrue(sets[:, 0].all())
-        self.assertFalse(sets[:, 1].all())
-        self.assertFalse(sets[:, 2].all())
+        self.assertFalse(sets[:, 1].any())
+        self.assertFalse(sets[:, 2].any())
 
     def test_no_threshold_raises(self):
         inferer = ConformalPredictor()  # qhat None
@@ -169,6 +188,47 @@ def net(x):
         with self.assertRaises(TypeError):
             inferer(torch.zeros(1, 1), net)
 
+    def test_set_threshold_rejects_non_scalar(self):
+        inferer = ConformalPredictor()
+        with self.assertRaises(ValueError):
+            inferer.set_threshold(torch.tensor([0.1, 0.2]))
+
+    def test_set_threshold_rejects_non_tensor(self):
+        inferer = ConformalPredictor()
+        with self.assertRaises(TypeError):
+            inferer.set_threshold(0.5)  # type: ignore[arg-type]
+
+    def test_calibrate_restores_training_state(self):
+        class Net(torch.nn.Module):
+            def forward(self, x):
+                logits = torch.zeros(x.shape[0], 2)
+                logits[:, 0] = 2.1972
+                return logits
+
+        net = Net()
+        net.train()
+        self.assertTrue(net.training)
+        cal_loader = [{"image": torch.zeros(2, 1), "label": torch.zeros(2, 1, dtype=torch.long)} for _ in range(5)]
+        inferer = ConformalPredictor(alpha=0.1)
+        inferer.calibrate(net, cal_loader, device=torch.device("cpu"))
+        # training state must be restored to True
+        self.assertTrue(net.training)
+
+    def test_calibrate_preserves_eval_state(self):
+        class Net(torch.nn.Module):
+            def forward(self, x):
+                logits = torch.zeros(x.shape[0], 2)
+                logits[:, 0] = 2.1972
+                return logits
+
+        net = Net()
+        net.eval()
+        self.assertFalse(net.training)
+        cal_loader = [{"image": torch.zeros(2, 1), "label": torch.zeros(2, 1, dtype=torch.long)} for _ in range(5)]
+        inferer = ConformalPredictor(alpha=0.1)
+        inferer.calibrate(net, cal_loader, device=torch.device("cpu"))
+        self.assertFalse(net.training)
+
 
 if __name__ == "__main__":
     unittest.main()