attempt to fix coverage

Author: fira7s

tests/datamodules/classification/test_imagenet.py

@@ -425,3 +425,74 @@ def test_get_indices_without_ood_or_shift(self, tmp_path):
         assert idx["near_oods"] == []
         assert idx["far_oods"] == []
         assert idx["shift"] == []
+
+    def test_tta_wraps_ood_sets(self, monkeypatch, tmp_path):
+        mod_name = ImageNetDataModule.__module__
+
+        monkeypatch.setattr(
+            f"{mod_name}.download_and_extract_hf_dataset",
+            self._fake_download_and_extract,
+            raising=True,
+        )
+        monkeypatch.setattr(
+            f"{mod_name}.download_and_extract_splits_from_hf",
+            self._fake_download_and_extract_splits_from_hf,
+            raising=True,
+        )
+        monkeypatch.setattr(
+            f"{mod_name}.FileListDataset",
+            self._DummyFileListDataset,
+            raising=True,
+        )
+        monkeypatch.setattr(
+            "torch_uncertainty.datamodules.classification.imagenet.get_ood_datasets",
+            self._fake_get_ood_datasets,
+        )
+
+        dm = ImageNetDataModule(
+            root=tmp_path,
+            batch_size=8,
+            eval_ood=True,
+            num_tta=2,
+            train_transform=nn.Identity(),
+            test_transform=nn.Identity(),
+            num_workers=0,
+            persistent_workers=False,
+            pin_memory=False,
+        )
+
+        dm.setup("test")
+
+        val_ood_wrapped = dm.get_val_ood_set()
+        test_ood_wrapped = dm.get_test_ood_set()
+        near_wrapped = dm.get_near_ood_set()
+        far_wrapped = dm.get_far_ood_set()
+
+        assert len(val_ood_wrapped) == len(dm.val_ood) * dm.num_tta
+        assert len(test_ood_wrapped) == len(dm.test_ood) * dm.num_tta
+        assert all(
+            len(w) == len(b) * dm.num_tta for w, b in zip(near_wrapped, dm.near_oods, strict=False)
+        )
+        assert all(
+            len(w) == len(b) * dm.num_tta for w, b in zip(far_wrapped, dm.far_oods, strict=False)
+        )
+
+        def _assert_first_block_repeat(wrapped_ds, num_tta: int):
+            if len(wrapped_ds) == 0 or num_tta < 2:
+                return
+            x0, y0 = wrapped_ds[0]
+            x1, y1 = wrapped_ds[1]
+            assert y0 == y1
+            if torch.is_tensor(x0) and torch.is_tensor(x1):
+                assert x0.shape == x1.shape
+            else:
+                assert type(x0) is type(x1)
+                if hasattr(x0, "size") and hasattr(x1, "size"):
+                    assert x0.size == x1.size
+
+        _assert_first_block_repeat(val_ood_wrapped, dm.num_tta)
+        _assert_first_block_repeat(test_ood_wrapped, dm.num_tta)
+        for w in near_wrapped:
+            _assert_first_block_repeat(w, dm.num_tta)
+        for w in far_wrapped:
+            _assert_first_block_repeat(w, dm.num_tta)

tests/routines/test_classification.py

@@ -1,3 +1,5 @@
+import logging
+import types
 from pathlib import Path
 
 import pytest
@@ -15,6 +17,7 @@
 from torch_uncertainty.losses import DECLoss, ELBOLoss
 from torch_uncertainty.ood.ood_criteria import (
     EntropyCriterion,
+    MaxSoftmaxCriterion,
 )
 from torch_uncertainty.post_processing import ConformalClsTHR
 from torch_uncertainty.routines import ClassificationRoutine
@@ -518,3 +521,137 @@ def patched_setup(self, stage=None):
         for needs_setup in {"react", "adascale_a", "vim", "knn", "nnguide"}:
             if crit == needs_setup:
                 assert getattr(c, "setup_flag", False), f"Setup not executed for '{crit}'."
+
+    def test_setup_logs_when_no_train_loader(self, caplog, monkeypatch):
+        dm = DummyClassificationDataModule(
+            root=Path(),
+            batch_size=4,
+            num_classes=3,
+            num_images=16,
+            eval_ood=True,
+        )
+
+        def _raise_train_loader(*_a, **_k):
+            raise RuntimeError("no train loader")
+
+        monkeypatch.setattr(
+            ClassificationRoutine, "_hyperparam_search_ood", lambda _self: None, raising=True
+        )
+        monkeypatch.setattr(dm, "train_dataloader", _raise_train_loader, raising=True)
+
+        model = dummy_ood_model(in_channels=3, feat_dim=64, num_classes=3)
+        routine = ClassificationRoutine(
+            model=model,
+            loss=None,
+            num_classes=3,
+            eval_ood=True,
+        )
+        routine.ood_criterion = MaxSoftmaxCriterion()  # no setup() side-effects
+
+        routine.trainer = types.SimpleNamespace(datamodule=dm)
+
+        with caplog.at_level(logging.INFO):
+            routine.setup("test")
+        assert any("No train loader detected" in r.message for r in caplog.records)
+
+    def test_create_near_far_metric_dicts_non_ensemble(self):
+        model = dummy_ood_model(in_channels=3, feat_dim=64, num_classes=3)
+        routine = ClassificationRoutine(
+            model=model, loss=None, num_classes=3, eval_ood=True, is_ensemble=False
+        )
+        routine.ood_criterion = MaxSoftmaxCriterion()
+
+        x = torch.rand(4, 3, 8, 8)
+        y = torch.tensor([0, 1, 2, 0])
+
+        class _DS:
+            def __init__(self, name):
+                self.dataset_name = name
+
+        routine.trainer = types.SimpleNamespace(
+            datamodule=types.SimpleNamespace(
+                get_indices=lambda: {"val_ood": 9, "near_oods": [2], "far_oods": [3], "shift": []},
+                near_oods=[_DS("nearX")],
+                far_oods=[_DS("farY")],
+            )
+        )
+
+        routine.test_step((x, y), batch_idx=0, dataloader_idx=2)  # near
+        assert "nearX" in routine.test_ood_metrics_near
+
+        routine.test_step((x, y), batch_idx=0, dataloader_idx=3)  # far
+        assert "farY" in routine.test_ood_metrics_far
+
+    def test_create_near_far_metric_dicts_ensemble_and_aggregator(self):
+        model = dummy_ood_model(in_channels=3, feat_dim=64, num_classes=3)
+        routine = ClassificationRoutine(
+            model=model, loss=None, num_classes=3, eval_ood=True, is_ensemble=True
+        )
+        routine.ood_criterion = MaxSoftmaxCriterion()
+
+        x = torch.rand(4, 3, 8, 8)
+        y = torch.tensor([0, 1, 2, 0])
+
+        class _DS:
+            def __init__(self, name):
+                self.dataset_name = name
+
+        routine.trainer = types.SimpleNamespace(
+            datamodule=types.SimpleNamespace(
+                get_indices=lambda: {
+                    "val_ood": 9,
+                    "near_oods": [5],
+                    "far_oods": [6],
+                    "shift": [7],
+                },
+                near_oods=[_DS("n1")],
+                far_oods=[_DS("f1")],
+            )
+        )
+
+        routine.test_step((x, y), batch_idx=0, dataloader_idx=1)  # aggregator
+        assert "n1" in routine.test_ood_ens_metrics_near
+        assert "f1" in routine.test_ood_ens_metrics_far
+
+        routine.test_step((x, y), batch_idx=0, dataloader_idx=5)  # near
+        routine.test_step((x, y), batch_idx=0, dataloader_idx=6)  # far
+        assert "n1" in routine.test_ood_ens_metrics_near
+        assert "f1" in routine.test_ood_ens_metrics_far
+
+    def test_skip_when_val_ood_loader(self):
+        model = dummy_ood_model(in_channels=3, feat_dim=64, num_classes=3)
+        routine = ClassificationRoutine(model=model, loss=None, num_classes=3, eval_ood=True)
+        routine.ood_criterion = MaxSoftmaxCriterion()
+
+        routine.trainer = types.SimpleNamespace(
+            datamodule=types.SimpleNamespace(
+                get_indices=lambda: {"val_ood": 4, "near_oods": [], "far_oods": [], "shift": []}
+            )
+        )
+        x = torch.rand(2, 3, 8, 8)
+        y = torch.tensor([0, 1])
+        routine.test_step((x, y), batch_idx=0, dataloader_idx=4)
+
+    def test_init_metrics_creates_shift_ens_metrics_when_ensemble_and_eval_shift(self):
+        model = dummy_ood_model(in_channels=3, feat_dim=64, num_classes=3)
+        routine = ClassificationRoutine(
+            model=model, loss=None, num_classes=3, eval_shift=True, is_ensemble=True
+        )
+        assert hasattr(routine, "test_shift_ens_metrics")
+
+    def test_shift_ens_update_path(self):
+        model = dummy_ood_model(in_channels=3, feat_dim=64, num_classes=3)
+        routine = ClassificationRoutine(
+            model=model, loss=None, num_classes=3, eval_shift=True, is_ensemble=True
+        )
+        routine.ood_criterion = MaxSoftmaxCriterion()
+
+        x = torch.rand(4, 3, 8, 8)
+        y = torch.tensor([0, 1, 2, 0])
+
+        routine.trainer = types.SimpleNamespace(
+            datamodule=types.SimpleNamespace(
+                get_indices=lambda: {"val_ood": 99, "near_oods": [], "far_oods": [], "shift": [7]}
+            )
+        )
+        routine.test_step((x, y), batch_idx=0, dataloader_idx=7)