👕 Merge pull request #109 from ENSTA-U2IS-AI/dev

👕 Small fixes and improvements

.github/workflows/run-tests.yml

@@ -70,7 +70,7 @@ jobs:
     - name: Check style & format
       if: steps.changed-files-specific.outputs.only_changed != 'true'
       run: |
-        python3 -m ruff check torch_uncertainty --no-fix
+        python3 -m ruff check torch_uncertainty --no-fix --statistics
         python3 -m ruff format torch_uncertainty --check
 
     - name: Test with pytest and compute coverage

.pre-commit-config.yaml

@@ -30,7 +30,7 @@ repos:
     language: python
     types_or: [python, pyi]
     require_serial: true
-    args: [--force-exclude, --fix, --exit-non-zero-on-fix]
+    args: [--force-exclude, --fix, --exit-non-zero-on-fix, --output-format, concise]
     exclude: ^auto_tutorials_source/
   - id: ruff-format
     name: ruff-format

README.md

@@ -26,7 +26,7 @@ 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:).
+- [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.
 
@@ -59,7 +59,7 @@ We also provide the following methods:
 
 ### Baselines
 
-To date, the following deep learning baselines have been implemented. **Click on the methods for tutorials**:
+To date, the following deep learning baselines have been implemented. **Click** :inbox_tray: **on the methods for tutorials**:
 
 - [Deep Ensembles](https://torch-uncertainty.github.io/auto_tutorials/tutorial_from_de_to_pe.html), BatchEnsemble, Masksembles, & MIMO
 - [MC-Dropout](https://torch-uncertainty.github.io/auto_tutorials/tutorial_mc_dropout.html)

auto_tutorials_source/tutorial_bayesian.py

@@ -37,8 +37,6 @@
 neural network utils from torch.nn, as well as the partial util to provide
 the modified default arguments for the ELBO loss.
 """
-
-# %%
 from pathlib import Path
 
 from lightning.pytorch import Trainer

auto_tutorials_source/tutorial_corruption.py

@@ -11,7 +11,6 @@
 torch_uncertainty.transforms.corruptions. We also need to load utilities from
 torchvision and matplotlib.
 """
-
 from torchvision.datasets import CIFAR10
 from torchvision.transforms import Compose, ToTensor, Resize
 

auto_tutorials_source/tutorial_evidential_classification.py

@@ -22,12 +22,8 @@
 - the evidential objective: the DECLoss from torch_uncertainty.losses
 - the datamodule that handles dataloaders & transforms: MNISTDataModule from torch_uncertainty.datamodules
 
-We also need to define an optimizer using torch.optim, the neural network utils within torch.nn, as well as the partial util to provide
-the modified default arguments for the DEC loss.
+We also need to define an optimizer using torch.optim, the neural network utils within torch.nn.
 """
-
-# %%
-from functools import partial
 from pathlib import Path
 
 import torch
@@ -73,11 +69,10 @@ def optim_lenet(model: nn.Module) -> dict:
 # %%
 # 4. The Loss and the Training Routine
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-# Next, we need to define the loss to be used during training. To do this, we
-# redefine the default parameters for the DEC loss using the partial
-# function from functools. After that, we define the training routine using
+# Next, we need to define the loss to be used during training.
+# After that, we define the training routine using
 # the single classification model training routine from
-# torch_uncertainty.routines.classification.ClassificationSingle.
+# torch_uncertainty.routines.ClassificationRoutine.
 # In this routine, we provide the model, the DEC loss, the optimizer,
 # and all the default arguments.
 
@@ -152,4 +147,4 @@ def rotated_mnist(angle: int) -> None:
 # References
 # ----------
 #
-# - **Deep Evidential Classification:** Murat Sensoy, Lance Kaplan, & Melih Kandemir (2018). Evidential Deep Learning to Quantify Classification Uncertainty `NeurIPS 2018 <https://arxiv.org/pdf/1806.01768>`_
+# - **Deep Evidential Classification:** Murat Sensoy, Lance Kaplan, & Melih Kandemir (2018). Evidential Deep Learning to Quantify Classification Uncertainty `NeurIPS 2018 <https://arxiv.org/pdf/1806.01768>`_.

auto_tutorials_source/tutorial_from_de_to_pe.py

@@ -1,11 +1,12 @@
-"""Improved Ensemble parameter-efficiency with Packed-Ensembles
+"""
+Improved Ensemble parameter-efficiency with Packed-Ensembles
 ============================================================
 
-*This tutorial is adapted from a notebook part of a lecture given at the [Helmholtz AI Conference](https://haicon24.de/) by Sebastian Starke, Peter Steinbach, Gianni Franchi, and Olivier Laurent.*
+*This tutorial is adapted from a notebook part of a lecture given at the `Helmholtz AI Conference <https://haicon24.de/>`_ by Sebastian Starke, Peter Steinbach, Gianni Franchi, and Olivier Laurent.*
 
 In this notebook will work on the MNIST dataset that was introduced by Corinna Cortes, Christopher J.C. Burges, and later modified by Yann LeCun in the foundational paper:
 
-- [Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. "Gradient-based learning applied to document recognition." Proceedings of the IEEE.](http://yann.lecun.com/exdb/publis/pdf/lecun-98.pdf)
+- `Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. "Gradient-based learning applied to document recognition." Proceedings of the IEEE. <http://yann.lecun.com/exdb/publis/pdf/lecun-98.pdf>`_.
 
 The MNIST dataset consists of 70 000 images of handwritten digits from 0 to 9. The images are grayscale and 28x28-pixel sized. The task is to classify the images into their respective digits. The dataset can be automatically downloaded using the `torchvision` library.
 
@@ -15,19 +16,19 @@
 - Calibration error: a measure of the calibration of the predicted probabilities,
 - Negative Log-Likelihood: the value of the loss on the test set.
 
-Throughout this notebook, we abstract the training and evaluation process using [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/)
-and [TorchUncertainty](https://torch-uncertainty.github.io/).
+Throughout this notebook, we abstract the training and evaluation process using `PyTorch Lightning <https://lightning.ai/docs/pytorch/stable/>`_
+and `TorchUncertainty <https://torch-uncertainty.github.io/>`_.
 
 Similarly to keras for tensorflow, PyTorch Lightning is a high-level interface for PyTorch that simplifies the training and evaluation process using a Trainer.
 TorchUncertainty is partly built on top of PyTorch Lightning and provides tools to train and evaluate models with uncertainty quantification.
 
 TorchUncertainty includes datamodules that handle the data loading and preprocessing. We don't use them here for tutorial purposes.
-"""
-# 1. Download, instantiate and visualize the datasets
-# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-#
-# The dataset is automatically downloaded using torchvision. We then visualize a few images to see a bit what we are working with.
 
+1. Download, instantiate and visualize the datasets
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The dataset is automatically downloaded using torchvision. We then visualize a few images to see a bit what we are working with.
+"""
 # Create the transforms for the images
 import torch
 import torchvision.transforms as T
@@ -138,7 +139,7 @@ def optim_recipe(model, lr_mult: float = 1.0):
 
 
 # %%
-# To train the model, we use [TorchUncertainty](https://torch-uncertainty.github.io/), a library that we have developed to ease
+# To train the model, we use `TorchUncertainty <https://torch-uncertainty.github.io/>`_, a library that we have developed to ease
 # the training and evaluation of models with uncertainty.
 #
 # **Note:** To train supervised classification models we most often use the cross-entropy loss.
@@ -240,7 +241,7 @@ def optim_recipe(model, lr_mult: float = 1.0):
 # We have put the pre-trained models on Hugging Face that you can download with the utility function
 # "hf_hub_download" imported just below. These models are trained for 75 epochs and are therefore not
 # comparable to the all the other models trained in this notebook. The pretrained models can be seen
-# [here](https://huggingface.co/ENSTA-U2IS/tutorial-models) and TorchUncertainty's are [here](https://huggingface.co/torch-uncertainty).
+# `here <https://huggingface.co/ENSTA-U2IS/tutorial-models>`_ and TorchUncertainty's are `here <https://huggingface.co/torch-uncertainty>`_.
 
 from torch_uncertainty.utils.hub import hf_hub_download
 
@@ -289,15 +290,15 @@ def optim_recipe(model, lr_mult: float = 1.0):
 # 4. From Deep Ensembles to Packed-Ensembles
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 #
-# In the paper [Packed-Ensembles for Efficient Uncertainty Quantification](https://arxiv.org/abs/2210.09184)
+# In the paper `Packed-Ensembles for Efficient Uncertainty Quantification <https://arxiv.org/abs/2210.09184>`_
 # published at the International Conference on Learning Representations (ICLR) in 2023, we introduced a
 # modification of Deep Ensembles to make it more computationally-efficient. The idea is to pack the ensemble
 # members into a single model, which allows us to train the ensemble in a single forward pass.
 # This modification is particularly useful when the ensemble size is large, as it is often the case in practice.
 #
 # We will need to update the model and replace the layers with their Packed equivalents. You can find the
-# documentation of the Packed-Linear layer [here](https://torch-uncertainty.github.io/generated/torch_uncertainty.layers.PackedLinear.html),
-# and the Packed-Conv2D, [here](https://torch-uncertainty.github.io/generated/torch_uncertainty.layers.PackedLinear.html).
+# documentation of the Packed-Linear layer `here <https://torch-uncertainty.github.io/generated/torch_uncertainty.layers.PackedLinear.html>`_,
+# and the Packed-Conv2D, `here <https://torch-uncertainty.github.io/generated/torch_uncertainty.layers.PackedLinear.html>`_.
 
 import torch
 import torch.nn as nn
@@ -387,7 +388,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 # %%
 # The training time should be approximately similar to the one of the single model that you trained before. However, please note that we are working with very small models, hence completely underusing your GPU. As such, the training time is not representative of what you would observe with larger models.
 #
-# You can read more on Packed-Ensembles in the [paper](https://arxiv.org/abs/2210.09184) or the [Medium](https://medium.com/@adrien.lafage/make-your-neural-networks-more-reliable-with-packed-ensembles-7ad0b737a873) post.
+# You can read more on Packed-Ensembles in the `paper <https://arxiv.org/abs/2210.09184>`_ or the `Medium <https://medium.com/@adrien.lafage/make-your-neural-networks-more-reliable-with-packed-ensembles-7ad0b737a873>`_ post.
 #
 # To Go Further & More Concepts of Uncertainty in ML
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -413,4 +414,4 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 # Grouping Loss
 # ^^^^^^^^^^^^^
 #
-# The grouping loss is a measure of uncertainty orthogonal to calibration. Have a look at [this paper](https://arxiv.org/abs/2210.16315) to learn about it. Check out their small library [GLest](https://github.com/aperezlebel/glest). TorchUncertainty includes a wrapper of the library to compute the grouping loss with eval_grouping_loss parameter.
+# The grouping loss is a measure of uncertainty orthogonal to calibration. Have a look at `this paper <https://arxiv.org/abs/2210.16315>`_ to learn about it. Check out their small library `GLest <https://github.com/aperezlebel/glest>`_. TorchUncertainty includes a wrapper of the library to compute the grouping loss with eval_grouping_loss parameter.

auto_tutorials_source/tutorial_mc_dropout.py

@@ -19,20 +19,19 @@
 
 First, we have to load the following utilities from TorchUncertainty:
 
-- the Trainer from Lightning
+- the TUTrainer from TorchUncertainty utils
 - the datamodule handling dataloaders: MNISTDataModule from torch_uncertainty.datamodules
-- the model: LeNet, which lies in torch_uncertainty.models
+- the model: lenet from torch_uncertainty.models
 - the MC Dropout wrapper: mc_dropout, from torch_uncertainty.models.wrappers
 - the classification training & evaluation routine in the torch_uncertainty.routines
 - an optimization recipe in the torch_uncertainty.optim_recipes module.
 
 We also need import the neural network utils within `torch.nn`.
 """
 
-# %%
 from pathlib import Path
 
-from lightning.pytorch import Trainer
+from torch_uncertainty.utils import TUTrainer
 from torch import nn
 
 from torch_uncertainty.datamodules import MNISTDataModule
@@ -42,18 +41,17 @@
 from torch_uncertainty.routines import ClassificationRoutine
 
 # %%
-# 2. Creating the necessary variables
-# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# 2. Defining the Model and the Trainer
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 #
-# In the following, we will need to define the root of the datasets and the
-# logs, and to fake-parse the arguments needed for using the PyTorch Lightning
-# Trainer. We also create the datamodule that handles the MNIST dataset,
+# In the following, we first create the trainer and instantiate 
+# the datamodule that handles the MNIST dataset,
 # dataloaders and transforms. We create the model using the
-# blueprint from torch_uncertainty.models and we wrap it into mc_dropout.
-#
-# It is important to add a ``dropout_rate`` argument in your model to use Monte Carlo dropout.
+# blueprint from torch_uncertainty.models and we wrap it into an mc_dropout.
+# To use the mc_dropout wrapper, **make sure that you use dropout modules** and
+# not functionals. Moreover, **they have to be** instantiated in the __init__ method.
 
-trainer = Trainer(accelerator="cpu", max_epochs=2, enable_progress_bar=False)
+trainer = TUTrainer(accelerator="cpu", max_epochs=2, enable_progress_bar=False)
 
 # datamodule
 root = Path("data")
@@ -63,18 +61,18 @@
 model = lenet(
     in_channels=datamodule.num_channels,
     num_classes=datamodule.num_classes,
-    dropout_rate=0.5,
+    dropout_rate=0.4,
 )
 
 mc_model = mc_dropout(model, num_estimators=16, last_layer=False)
 
 # %%
 # 3. The Loss and the Training Routine
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-# This is a classification problem, and we use CrossEntropyLoss as the likelihood.
+# This is a classification problem, and we use CrossEntropyLoss as the (negative-log-)likelihood.
 # We define the training routine using the classification training routine from
-# torch_uncertainty.routines.classification. We provide the number of classes
-# and channels, the optimizer wrapper, and the dropout rate.
+# torch_uncertainty.routines. We provide the number of classes
+# the optimization recipe and tell the routine that our model is an ensemble at evaluation time.
 
 routine = ClassificationRoutine(
     num_classes=datamodule.num_classes,
@@ -87,15 +85,19 @@
 # %%
 # 4. Gathering Everything and Training the Model
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# We can now train the model using the trainer. We pass the routine and the datamodule
+# to the fit and test methods of the trainer. It will automatically evaluate some uncertainty
+# metrics that you will find in the table below.
 
 trainer.fit(model=routine, datamodule=datamodule)
-trainer.test(model=routine, datamodule=datamodule)
+results = trainer.test(model=routine, datamodule=datamodule)
 
 # %%
 # 5. Testing the Model
 # ~~~~~~~~~~~~~~~~~~~~
 # Now that the model is trained, let's test it on MNIST. Don't forget to call
-# .eval() to enable dropout at inference.
+# .eval() to enable dropout at evaluation and get multiple (here 16) predictions.
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -127,7 +129,7 @@ def imshow(img):
 for j in range(6):
     values, predicted = torch.max(probs[:, j], 1)
     print(
-        f"Predicted digits for the image {j+1}: ",
+        f"MC-Dropout predictions for the image {j+1}: ",
         " ".join([str(image_id.item()) for image_id in predicted]),
     )
 

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.2.1"
+release = "0.2.1.post0"
 
 # -- 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.2.1"
+version = "0.2.1.post0"
 authors = [
     { name = "ENSTA U2IS", email = "olivier.laurent@ensta-paris.fr" },
     { name = "Adrien Lafage", email = "adrienlafage@outlook.com" },
@@ -51,7 +51,7 @@ image = ["scikit-image", "h5py"]
 tabular = ["pandas"]
 dev = [
     "torch_uncertainty[image]",
-    "ruff==0.3.4",
+    "ruff==0.5.2",
     "pytest-cov",
     "pre-commit",
     "pre-commit-hooks",
@@ -93,6 +93,7 @@ lint.extend-select = [
     "ISC",
     "ICN",
     "N",
+    "NPY",
     "PERF",
     "PIE",
     "PTH",

tests/_dummies/dataset.py

@@ -53,7 +53,7 @@ def __init__(
         else:
             shape = (num_images, num_channels, image_size, image_size)
 
-        self.data = np.random.randint(
+        self.data = np.random.default_rng().integers(
             low=0,
             high=255,
             size=shape,
@@ -187,14 +187,15 @@ def __init__(
 
         smnt_shape = (num_images, 1, image_size, image_size)
 
-        self.data = np.random.randint(
+        rng = np.random.default_rng()
+        self.data = rng.integers(
             low=0,
             high=255,
             size=img_shape,
             dtype=np.uint8,
         )
 
-        self.targets = np.random.randint(
+        self.targets = rng.integers(
             low=0,
             high=num_classes,
             size=smnt_shape,
@@ -245,20 +246,18 @@ def __init__(
         else:
             smnt_shape = (num_images, output_dim, image_size, image_size)
 
-        self.data = np.random.randint(
+        rng = np.random.default_rng()
+        self.data = rng.integers(
             low=0,
             high=255,
             size=img_shape,
             dtype=np.uint8,
         )
 
-        self.targets = (
-            np.random.uniform(
-                low=0,
-                high=1,
-                size=smnt_shape,
-            )
-            * 100
+        self.targets = rng.uniform(
+            low=0,
+            high=100,
+            size=smnt_shape,
         )
 
     def __getitem__(self, index: int) -> tuple[Any, Any]:

tests/metrics/classification/test_risk_coverage.py

@@ -52,6 +52,9 @@ def test_plot(self) -> None:
         assert ax.get_ylabel() == "Risk - Error Rate (%)"
         plt.close(fig)
 
+        metric = AURC()
+        assert metric(torch.zeros(1), torch.zeros(1)).isnan()
+
         metric = AURC()
         metric.update(scores, values)
         fig, ax = metric.plot(plot_value=False)
@@ -91,6 +94,9 @@ def test_compute_zero(self) -> None:
         metric = CovAtxRisk(risk_threshold=0.5)
         assert metric(probs, targets) == 1
 
+        metric = CovAtxRisk(risk_threshold=0.5)
+        assert metric(torch.zeros(0), torch.zeros(0)).isnan()
+
     def test_errors(self):
         with pytest.raises(
             TypeError, match="Expected threshold to be of type float"