🔨 Rework PPs, Use transformsV2, & Fix corruptions

.github/workflows/run-tests.yml

@@ -36,23 +36,23 @@ jobs:
         echo "PYTHON_VERSION=$(python -c "import platform; print(platform.python_version())")"
         echo "PYTHON_VERSION=$(python -c "import platform; print(platform.python_version())")" >> $GITHUB_ENV
 
-    - name: Get changed files
-      id: changed-files-specific
-      uses: tj-actions/changed-files@v42
-      with:
-        files: |
-          auto_tutorials_source/**
-          data/**
-          experiments/**
-          docs/**
-          *.md
-          *.yaml
-          *.yml
-          LICENSE
-          .gitignore
+    # - name: Get changed files
+    #   id: changed-files-specific
+    #   uses: tj-actions/changed-files@v42
+    #   with:
+    #     files: |
+    #       auto_tutorials_source/**
+    #       data/**
+    #       experiments/**
+    #       docs/**
+    #       *.md
+    #       *.yaml
+    #       *.yml
+    #       LICENSE
+    #       .gitignore
 
     - name: Cache folder for TorchUncertainty
-      if: steps.changed-files-specific.outputs.only_changed != 'true'
+      # if: steps.changed-files-specific.outputs.only_changed != 'true'
       uses: actions/cache@v4
       id: cache-folder
       with:
@@ -61,43 +61,45 @@ jobs:
         key: torch-uncertainty-${{ runner.os }}
 
     - name: Install dependencies
-      if: steps.changed-files-specific.outputs.only_changed != 'true'
+      # if: steps.changed-files-specific.outputs.only_changed != 'true'
       run: |
         python3 -m pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu
         python3 -m pip install .[all]
 
     - name: Check style & format
-      if: steps.changed-files-specific.outputs.only_changed != 'true'
+      # if: steps.changed-files-specific.outputs.only_changed != 'true'
       run: |
         python3 -m ruff check torch_uncertainty --no-fix --statistics
         python3 -m ruff format torch_uncertainty --check
 
     - name: Test with pytest and compute coverage
-      if: steps.changed-files-specific.outputs.only_changed != 'true'
+      # if: steps.changed-files-specific.outputs.only_changed != 'true'
       run: |
         python3 -m pytest --cov --cov-report xml --durations 10 --junitxml=junit.xml
 
     - name: Upload coverage to Codecov
-      if: steps.changed-files-specific.outputs.only_changed != 'true' && (github.event_name != 'pull_request' || github.base_ref == 'dev')
+      # if: steps.changed-files-specific.outputs.only_changed != 'true' && (github.event_name != 'pull_request' || github.base_ref == 'dev')
+      if: github.event_name != 'pull_request' || github.base_ref == 'dev'
       uses: codecov/codecov-action@v4
       continue-on-error: true
       with:
         token: ${{ secrets.CODECOV_TOKEN }}
         file: coverage.xml
-        flags: cpu,pytest
+        flags: pytest
         name: CPU-coverage
         env_vars: PYTHON_VERSION
 
     - name: Upload test results to Codecov
-      if: steps.changed-files-specific.outputs.only_changed != 'true' && (github.event_name != 'pull_request' || github.base_ref == 'dev')
+      # if: steps.changed-files-specific.outputs.only_changed != 'true' && (github.event_name != 'pull_request' || github.base_ref == 'dev')
+      if: github.event_name != 'pull_request' || github.base_ref == 'dev'
       uses: codecov/test-results-action@v1
       continue-on-error: true
       with:
         token: ${{ secrets.CODECOV_TOKEN }}
-        flags: cpu,pytest
+        flags: pytest
         env_vars: PYTHON_VERSION
 
     - name: Test sphinx build without tutorials
-      if: steps.changed-files-specific.outputs.only_changed != 'true'
+      # if: steps.changed-files-specific.outputs.only_changed != 'true'
       run: |
         cd docs && make clean && make html-noplot

README.md

@@ -26,9 +26,9 @@ This package provides a multi-level API, including:
 
 - easy-to-use :zap: lightning **uncertainty-aware** training & evaluation routines for **4 tasks**: classification, probabilistic and pointwise regression, and segmentation.
 - ready-to-train baselines on research datasets, such as ImageNet and CIFAR
-- [pretrained weights](https://huggingface.co/torch-uncertainty) for these baselines on ImageNet and CIFAR ( :construction: work in progress :construction: ).
 - **layers**, **models**, **metrics**, & **losses** available for use in your networks
 - scikit-learn style post-processing methods such as Temperature Scaling.
+- transformations, including corruptions resulting in additional "corrupted datasets" available on [HuggingFace](https://huggingface.co/torch-uncertainty)
 
 Have a look at the [Reference page](https://torch-uncertainty.github.io/references.html) or the [API reference](https://torch-uncertainty.github.io/api.html) for a more exhaustive list of the implemented methods, datasets, metrics, etc.
 

auto_tutorials_source/tutorial_mc_batch_norm.py

@@ -86,13 +86,13 @@
 # MCBatchNorm layers, and that we want to use 8 stochastic estimators.
 # The amount of stochasticity is controlled by the ``mc_batch_size`` argument.
 # The larger the ``mc_batch_size``, the more stochastic the predictions will be.
-# The authors suggest 32 as a good value for ``mc_batch_size`` but we use 4 here
+# The authors suggest 32 as a good value for ``mc_batch_size`` but we use 16 here
 # to highlight the effect of stochasticity on the predictions.
 
 routine.model = MCBatchNorm(
     routine.model, num_estimators=8, convert=True, mc_batch_size=16
 )
-routine.model.fit(datamodule.train)
+routine.model.fit(dataloader=datamodule.postprocess_dataloader())
 routine = routine.eval()  # To avoid prints
 
 # %%

auto_tutorials_source/tutorial_scaler.py

@@ -6,8 +6,8 @@
 of the top-label predictions and the reliability of the underlying neural network.
 
 This tutorial provides extensive details on how to use the TemperatureScaler
-class, however, this is done automatically in the classification routine when setting
-the `calibration_set` to val or test.
+class, however, this is done automatically in the datamodule when setting
+the `postprocess_set` to val or test.
 
 Through this tutorial, we also see how to use the datamodules outside any Lightning trainers,
 and how to use TorchUncertainty's models.
@@ -57,12 +57,12 @@
 # element if eval_ood is True: the dataloader of in-distribution data and the dataloader
 # of out-of-distribution data. Otherwise, it is a list of 1 element.
 
-dm = CIFAR100DataModule(root="./data", eval_ood=False, batch_size=32)
+dm = CIFAR100DataModule(root="./data", eval_ood=False, batch_size=32, postprocess_set="test")
 dm.prepare_data()
 dm.setup("test")
 
-# Get the full test dataloader (unused in this tutorial)
-dataloader = dm.test_dataloader()[0]
+# Get the full post-processing dataloader (unused in this tutorial)
+dataloader = dm.postprocess_dataloader()
 
 # %%
 # 4. Iterating on the Dataloader and Computing the ECE
@@ -84,6 +84,7 @@
     dataset, [1000, 1000, len(dataset) - 2000]
 )
 test_dataloader = DataLoader(test_dataset, batch_size=32)
+calibration_dataloader = DataLoader(cal_dataset, batch_size=32)
 
 # Initialize the ECE
 ece = CalibrationError(task="multiclass", num_classes=100)
@@ -114,7 +115,7 @@
 
 # Fit the scaler on the calibration dataset
 scaled_model = TemperatureScaler(model=model)
-scaled_model.fit(calibration_set=cal_dataset)
+scaled_model.fit(dataloader=calibration_dataloader)
 
 # %%
 # 6. Iterating Again to Compute the Improved ECE

docs/source/conf.py

@@ -15,7 +15,7 @@
     f"{datetime.now().year!s}, Adrien Lafage and Olivier Laurent"
 )
 author = "Adrien Lafage and Olivier Laurent"
-release = "0.4.2"
+release = "0.4.3"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "torch_uncertainty"
-version = "0.4.2"
+version = "0.4.3"
 authors = [
     { name = "ENSTA U2IS", email = "olivier.laurent@ensta-paris.fr" },
     { name = "Adrien Lafage", email = "adrienlafage@outlook.com" },
@@ -46,7 +46,6 @@ experiments = [
     "safetensors",
 ]
 image = [
-    "scikit-image", 
     "kornia",
     "h5py",
     "opencv-python",

tests/_dummies/datamodule.py

@@ -2,10 +2,10 @@
 
 import numpy as np
 import torch
-import torchvision.transforms.v2 as T
 from numpy.typing import ArrayLike
 from torch.utils.data import DataLoader
 from torchvision import tv_tensors
+from torchvision.transforms import v2
 
 from torch_uncertainty.datamodules import TUDataModule
 
@@ -52,8 +52,8 @@ def __init__(
         self.ood_dataset = DummyClassificationDataset
         self.shift_dataset = DummyClassificationDataset
 
-        self.train_transform = T.ToTensor()
-        self.test_transform = T.ToTensor()
+        self.train_transform = v2.ToTensor()
+        self.test_transform = v2.ToTensor()
 
     def prepare_data(self) -> None:
         pass
@@ -207,15 +207,15 @@ def __init__(
 
         self.dataset = DummySegmentationDataset
 
-        self.train_transform = T.ToDtype(
+        self.train_transform = v2.ToDtype(
             dtype={
                 tv_tensors.Image: torch.float32,
                 tv_tensors.Mask: torch.int64,
                 "others": None,
             },
             scale=True,
         )
-        self.test_transform = T.ToDtype(
+        self.test_transform = v2.ToDtype(
             dtype={
                 tv_tensors.Image: torch.float32,
                 tv_tensors.Mask: torch.int64,
@@ -296,15 +296,15 @@ def __init__(
 
         self.dataset = DummPixelRegressionDataset
 
-        self.train_transform = T.ToDtype(
+        self.train_transform = v2.ToDtype(
             dtype={
                 tv_tensors.Image: torch.float32,
                 tv_tensors.Mask: torch.float32,
                 "others": None,
             },
             scale=True,
         )
-        self.test_transform = T.ToDtype(
+        self.test_transform = v2.ToDtype(
             dtype={
                 tv_tensors.Image: torch.float32,
                 tv_tensors.Mask: torch.float32,

tests/datamodules/classification/test_cifar10.py

@@ -10,7 +10,7 @@ class TestCIFAR10DataModule:
     """Testing the CIFAR10DataModule datamodule class."""
 
     def test_cifar10_main(self):
-        dm = CIFAR10DataModule(root="./data/", batch_size=128, cutout=16)
+        dm = CIFAR10DataModule(root="./data/", batch_size=128, cutout=16, postprocess_set="test")
 
         assert dm.dataset == CIFAR10
         assert isinstance(dm.train_transform.transforms[2], Cutout)

tests/datamodules/segmentation/test_muad.py

@@ -35,7 +35,3 @@ def test_camvid_main(self):
         dm.setup()
         dm.train_dataloader()
         dm.val_dataloader()
-
-    def test_small_muad_accessibility(self):
-        dataset = MUAD(root="./data/", split="test", version="small", download=True)
-        assert len(dataset.samples) > 0, "Dataset is not found"

tests/datamodules/test_abstract_datamodule.py

@@ -47,12 +47,29 @@ def test_cv_main(self):
 
     def test_errors(self):
         TUDataModule.__abstractmethods__ = set()
-        dm = TUDataModule("root", 128, 0.0, 4, True, True)
+        dm = TUDataModule(
+            root="root",
+            batch_size=128,
+            val_split=0.0,
+            num_workers=4,
+            pin_memory=True,
+            persistent_workers=True,
+        )
         ds = DummyClassificationDataset(Path("root"))
         dm.train = ds
         dm.val = ds
         dm.test = ds
-        cv_dm = CrossValDataModule("root", [0], [1], dm, 128, 0.0, 4, True, True)
+        cv_dm = CrossValDataModule(
+            root="root",
+            train_idx=[0],
+            val_idx=[1],
+            datamodule=dm,
+            batch_size=128,
+            val_split=0.0,
+            num_workers=4,
+            pin_memory=True,
+            persistent_workers=True,
+        )
         with pytest.raises(NotImplementedError):
             cv_dm.setup()
             cv_dm._get_train_data()

tests/post_processing/test_laplace.py

@@ -1,6 +1,6 @@
 import torch
 from torch import nn
-from torch.utils.data import TensorDataset
+from torch.utils.data import DataLoader, TensorDataset
 
 from tests._dummies.model import dummy_model
 from torch_uncertainty.post_processing import LaplaceApprox, PostProcessing
@@ -20,12 +20,12 @@ class TestLaplace:
     """Testing the LaplaceApprox class."""
 
     def test_training(self):
-        ds = TensorDataset(torch.randn(16, 1), torch.randn(16, 10))
+        dl = DataLoader(TensorDataset(torch.randn(16, 1), torch.randn(16, 10)), batch_size=5)
         la = LaplaceApprox(
             task="classification",
             model=dummy_model(1, 10),
         )
-        la.fit(ds)
+        la.fit(dl)
         la(torch.randn(1, 1))
         la = LaplaceApprox(task="classification")
         la.set_model(dummy_model(1, 10))

tests/post_processing/test_mc_batch_norm.py

@@ -4,6 +4,7 @@
 import torch
 import torchvision.transforms as T
 from torch import nn
+from torch.utils.data import DataLoader
 
 from tests._dummies.dataset import DummyClassificationDataset
 from torch_uncertainty.layers.mc_batch_norm import MCBatchNorm2d
@@ -21,10 +22,7 @@ def test_main(self):
 
         model = lenet(1, 1, norm=nn.BatchNorm2d)
         stoch_model = MCBatchNorm(
-            nn.Sequential(model),
-            num_estimators=2,
-            convert=True,
-            mc_batch_size=1,
+            nn.Sequential(model), num_estimators=2, convert=True, mc_batch_size=1
         )
         dataset = DummyClassificationDataset(
             "./",
@@ -34,7 +32,7 @@ def test_main(self):
             num_images=2,
             transform=T.ToTensor(),
         )
-        stoch_model.fit(dataset=dataset)
+        stoch_model.fit(dataloader=DataLoader(dataset, batch_size=6, shuffle=True))
         stoch_model.train()
         stoch_model(torch.randn(1, 1, 20, 20))
         stoch_model.eval()
@@ -48,8 +46,6 @@ def test_errors(self):
         model = nn.Identity()
         with pytest.raises(ValueError):
             MCBatchNorm(model, num_estimators=0, convert=True)
-        with pytest.raises(ValueError, match="mc_batch_size must be a positive integer"):
-            MCBatchNorm(model, num_estimators=1, convert=True, mc_batch_size=-1)
         with pytest.raises(ValueError):
             MCBatchNorm(model, num_estimators=1, convert=False)
         with pytest.raises(ValueError):
@@ -64,9 +60,10 @@ def test_errors(self):
             num_images=2,
             transform=T.ToTensor(),
         )
+        dl = DataLoader(dataset, batch_size=2, shuffle=True)
         stoch_model.eval()
         with pytest.raises(RuntimeError):
             stoch_model(torch.randn(1, 1, 20, 20))
 
         with pytest.raises(ValueError):
-            stoch_model.fit(dataset=dataset)
+            stoch_model.fit(dataloader=dl)