TorchUncertainty

https://torch-uncertainty.github.io/
Admin
TorchUncertainty is a practitioner-friendly PyTorch library for uncertainty quantification in...
Tokens:481,195
Snippets:2,683
Trust Score:9.7
Update:6 months ago
Show doc for...
Tokens:
Raw
### Install torch-uncertainty with all dependencies

Source: https://torch-uncertainty.github.io/installation.html

Install the package from source with all available optional dependencies.

```bash
pip install -e .[all]
```

--------------------------------

### setup

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datamodules/classification/mnist.html

Configures the datasets for training, validation, and testing stages.

```APIDOC
## setup

### Description
Sets up the datasets for training, validation, and testing based on the specified stage.

### Parameters
#### Arguments
- **stage** (Literal["fit", "test"] | None) - Optional - Stage of the setup process. Defaults to None.
```

--------------------------------

### Install Torch-Uncertainty with CPU Support

Source: https://torch-uncertainty.github.io/_sources/contributing.rst.txt

Use this command to install torch-uncertainty with development packages for systems without GPUs. Ensure uv is installed first.

```bash
uv sync --extra cpu
```

--------------------------------

### Install with Optional Dependencies

Source: https://torch-uncertainty.github.io/_sources/installation.rst.txt

Commands to install the package with specific dependency groups for development, documentation, or full functionality.

```bash
pip install -e .[dev,docs]
```

```bash
pip install -e .[all]
```

--------------------------------

### Setup TorchUncertainty Trainer and Routine

Source: https://torch-uncertainty.github.io/auto_tutorials/Ensemble_Methods/tutorial_from_de_to_pe.html

Initializes the TUTrainer for managing the training process and ClassificationRoutine for handling classification-specific logic, metrics, and loss. Ensure PyTorch and TorchUncertainty are installed.

```python
from torch_uncertainty import TUTrainer
from torch_uncertainty.routines import ClassificationRoutine

# Create the trainer that will handle the training
trainer = TUTrainer(accelerator="gpu", devices=1, max_epochs=MAX_EPOCHS, enable_progress_bar=False)

# The routine is a wrapper of the model that contains the training logic with the metrics, etc
routine = ClassificationRoutine(
    num_classes=10,
    model=model,
    loss=nn.CrossEntropyLoss(),
    optim_recipe=optim_recipe(model),
    eval_ood=True,
)
```

--------------------------------

### Install Torch Uncertainty via PyPI

Source: https://torch-uncertainty.github.io/_sources/installation.rst.txt

Standard installation command for the package using pip.

```bash
pip install torch-uncertainty
```

--------------------------------

### Setup Datasets

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datamodules/classification/mnist.html

Initializes datasets for training, validation, and testing stages, including optional train/val splits.

```python
    def setup(self, stage: Literal["fit", "test"] | None = None) -> None:
        """Set up the datasets for training, validation, and testing.

        Args:
            stage (Literal["fit", "test"] | None): Stage of the setup process. Defaults to ``None``.
        """
        if stage == "fit" or stage is None:
            full = self.dataset(
                self.root,
                train=True,
                download=False,
                transform=self.train_transform,
            )
            if self.val_split:
                self.train, self.val = create_train_val_split(
                    full,
                    self.val_split,
                    self.test_transform,
                )
            else:
                self.train = full
                self.val = self.dataset(
                    self.root,
                    train=False,
                    download=False,
                    transform=self.test_transform,
                )
        if stage == "test" or stage is None:
            self.test = self.dataset(
                self.root,
                train=False,
                download=False,
                transform=self.test_transform,
            )
        if stage not in ["fit", "test", None]:
            raise ValueError(f"Stage {stage} is not supported.")

        if self.eval_ood:
            self.ood = self.ood_dataset(
                self.root,
                download=False,
                transform=self.ood_transform,
            )
        if self.eval_shift:
            self.shift = self.shift_dataset(
                self.root,
                download=False,
                transform=self.test_transform,
            )
```

--------------------------------

### Install Torch-Uncertainty with GPU Support

Source: https://torch-uncertainty.github.io/_sources/contributing.rst.txt

Use this command to install torch-uncertainty with development packages for GPU-based systems. Ensure uv is installed first.

```bash
uv sync --extra gpu
```

--------------------------------

### Import Libraries and Setup Environment

Source: https://torch-uncertainty.github.io/auto_tutorials/Classification/tutorial_distribution_shift.html

Import necessary PyTorch and TorchUncertainty components, including the trainer, datamodules, models, and routines. This setup is essential for subsequent data handling and model training.

```python
from torch import nn, optim

from torch_uncertainty import TUTrainer
from torch_uncertainty.datamodules import CIFAR10DataModule
from torch_uncertainty.models.classification.resnet import resnet
from torch_uncertainty.routines.classification import ClassificationRoutine

```

--------------------------------

### Configure Data Module Setup

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datamodules/classification/imagenet.html

Handles dataset splitting and OOD/shift dataset initialization for fit and test stages.

```python
    def setup(self, stage: Literal["fit", "test"] | None = None) -> None:
        if stage == "fit" or stage is None:
            if self.test_alt is not None:
                raise ValueError("The test_alt argument is not supported for training.")
            full = self.dataset(
                self.root,
                split="train",
                transform=self.train_transform,
            )
            if self.val_split and isinstance(self.val_split, float):
                self.train, self.val = create_train_val_split(
                    full,
                    self.val_split,
                    self.test_transform,
                )
            elif isinstance(self.val_split, Path):
                self.train = Subset(full, self.train_indices)
                # TODO: improve the performance
                self.val = copy.deepcopy(Subset(full, self.val_indices))
                self.val.dataset.transform = self.test_transform
            else:
                self.train = full
                self.val = self.dataset(
                    self.root,
                    split="val",
                    transform=self.test_transform,
                )
        if stage == "test" or stage is None:
            self.test = self.dataset(
                self.root,
                split="val",
                transform=self.test_transform,
            )
        if stage not in ["fit", "test", None]:
            raise ValueError(f"Stage {stage} is not supported.")

        if self.eval_ood:
            if self.ood_ds == "inaturalist":
                self.ood = self.ood_dataset(
                    self.root,
                    version="2021_valid",
                    transform=self.test_transform,
                )
            else:
                self.ood = self.ood_dataset(
                    self.root,
                    transform=self.test_transform,
                    download=True,
                )

        if self.eval_shift:
            self.shift = self.shift_dataset(
                self.root,
                download=False,
                transform=self.test_transform,
                shift_severity=self.shift_severity,
            )
```

--------------------------------

### Install torch-uncertainty with dev and docs dependencies

Source: https://torch-uncertainty.github.io/installation.html

Install the package from source with additional dependencies required for development and documentation. This is necessary for contributing to the project.

```bash
pip install -e .[dev,docs]
```

--------------------------------

### Implement Data Setup and Splitting

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datamodules/segmentation/cityscapes.html

Handles the loading of datasets and creation of training/validation splits based on the provided stage.

```python
def setup(self, stage: str | None = None) -> None:
        if stage == "fit" or stage is None:
            full = self.dataset(
                root=self.root,
                split="train",
                mode=self.mode,
                target_type="semantic",
                transforms=self.train_transform,
            )

            if self.val_split is not None:
                self.train, self.val = create_train_val_split(
                    full,
                    self.val_split,
                    self.test_transform,
                )
            else:
                self.train = full
                self.val = self.dataset(
                    root=self.root,
                    split="val",
                    mode=self.mode,
                    target_type="semantic",
                    transforms=self.test_transform,
                )

        if stage == "test" or stage is None:
            self.test = self.dataset(
                root=self.root,
                split="val",
                mode=self.mode,
                target_type="semantic",
                transforms=self.test_transform,
            )

        if stage not in ["fit", "test", None]:
            raise ValueError(f"Stage {stage} is not supported.")
```

--------------------------------

### MCBatchNorm Initialization and Setup

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/post_processing/mc_batch_norm.html

Details on initializing and setting up the MCBatchNorm wrapper, including model conversion and parameter configuration.

```APIDOC
## MCBatchNorm

### Description
Monte Carlo Batch Normalization wrapper for 2d inputs. This class enhances a given PyTorch model by replacing standard BatchNorm2d layers with MCBatchNorm2d layers, enabling stochastic behavior in evaluation mode for uncertainty estimation.

### Method
`__init__`

### Parameters
#### Arguments
- **model** (`nn.Module` | `None`) - The PyTorch model to be converted. Defaults to `None`.
- **num_estimators** (`int`) - The number of Monte Carlo estimators to use. Defaults to `16`.
- **convert** (`bool`) - Whether to automatically convert BatchNorm2d layers to MCBatchNorm2d. Defaults to `True`.
- **mc_batch_size** (`int`) - The batch size to use for Monte Carlo estimation during fitting. Defaults to `32`.
- **device** (`Literal["cpu", "cuda"]` | `torch.device` | `None`) - The device on which to run the model. Defaults to `None`.

### Warning
The update of batch statistics in this implementation differs slightly from the original paper's wording. Instead of updating training-based statistics with one new batch, a direct replacement is performed. Refer to the issue <https://github.com/torch-uncertainty/torch-uncertainty/issues/218> for details.

### Note
This wrapper will be stochastic only in evaluation mode.

### References
[1] Teye M, Azizpour H, Smith K. Bayesian uncertainty estimation for batch normalized deep networks. In ICML 2018 <https://arxiv.org/abs/1802.06455>
```

--------------------------------

### Example Output

Source: https://torch-uncertainty.github.io/auto_tutorials/Classification/tutorial_pe_cifar10.html

Sample output showing the labels of the displayed images.

```text
dog   truck dog   deer
```

--------------------------------

### Clone Repository and Setup Environment

Source: https://torch-uncertainty.github.io/_sources/installation.rst.txt

Steps to clone the source code and prepare a dedicated conda environment.

```bash
git clone https://github.com/torch-uncertainty/torch-uncertainty.git
cd torch-uncertainty
```

```bash
conda create -n uncertainty python=3.12
conda activate uncertainty
```

--------------------------------

### Configure and Run Ensemble Routine

Source: https://torch-uncertainty.github.io/_sources/auto_tutorials/Ensemble_Methods/tutorial_from_de_to_pe.rst.txt

Setup the trainer and execute the test routine for an ensemble model.

```python
        format_batch_fn=RepeatTarget(8),  # How to handle the targets when comparing the predictions
        optim_recipe=None,  # No optim recipe as the model is already trained
        eval_ood=True,  # We want to evaluate the OOD-related metrics
    )

    trainer = TUTrainer(accelerator="gpu", devices=1, max_epochs=MAX_EPOCHS)

    ens_perf = trainer.test(ens_routine, dataloaders=[test_dl, ood_dl])
```

--------------------------------

### Install Pre-commit Hooks

Source: https://torch-uncertainty.github.io/_sources/contributing.rst.txt

Install pre-commit hooks to ensure code formatting and linting before committing. This command should be run from the project root.

```bash
uv run pre-commit install
```

--------------------------------

### Handle Training Start

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/routines/pixel_regression.html

Logs hyperparameters to the logger if available and validates the loss function.

```python
    def on_train_start(self) -> None:  # coverage: ignore
        """Put the hyperparameters in tensorboard."""
        if self.loss is None:
            raise ValueError(
                "To train a model, you must specify the `loss` argument in the routine. Got None."
            )
        if self.logger is not None:
            self.logger.log_hyperparams(
                self.hparams,
            )
```

--------------------------------

### Handle Test Start

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/routines/pixel_regression.html

Updates batch normalization layers and calculates parameter count if not already set.

```python
    def on_test_start(self) -> None:
        """Prepare the test step.

        Update the batchnorms if needed.
        """
        if hasattr(self.model, "need_bn_update"):
            self.model.bn_update(self.trainer.train_dataloader, device=self.device)

        if self.num_params is None:
            self.num_params = sum(p.numel() for p in self.model.parameters())
```

--------------------------------

### Training Output Logs

Source: https://torch-uncertainty.github.io/_sources/auto_tutorials/Bayesian_Methods/tutorial_bayesian.rst.txt

Example output showing progress bars and model summary statistics during the training process.

```none
  0%|          | 0.00/9.91M [00:00<?, ?B/s]
  1%|          | 98.3k/9.91M [00:00<00:15, 616kB/s]
  4%|▍         | 393k/9.91M [00:00<00:07, 1.34MB/s]
  6%|▋         | 623k/9.91M [00:00<00:06, 1.37MB/s]
 16%|█▌        | 1.54M/9.91M [00:00<00:02, 3.09MB/s]
 24%|██▍       | 2.39M/9.91M [00:00<00:01, 3.88MB/s]
 30%|███       | 3.01M/9.91M [00:00<00:01, 4.20MB/s]
 35%|███▌      | 3.47M/9.91M [00:01<00:01, 4.09MB/s]
 40%|████      | 4.00M/9.91M [00:01<00:01, 3.84MB/s]
 47%|████▋     | 4.65M/9.91M [00:01<00:01, 3.92MB/s]
 54%|█████▍    | 5.34M/9.91M [00:01<00:01, 4.03MB/s]
 61%|██████    | 6.03M/9.91M [00:01<00:00, 4.12MB/s]
 68%|██████▊   | 6.72M/9.91M [00:01<00:00, 4.17MB/s]
 75%|███████▍  | 7.41M/9.91M [00:02<00:00, 4.20MB/s]
 82%|████████▏ | 8.13M/9.91M [00:02<00:00, 4.29MB/s]
 89%|████████▉ | 8.81M/9.91M [00:02<00:00, 4.29MB/s]
 95%|█████████▍| 9.37M/9.91M [00:02<00:00, 4.03MB/s]
100%|██████████| 9.91M/9.91M [00:02<00:00, 3.85MB/s]

  0%|          | 0.00/28.9k [00:00<?, ?B/s]
100%|██████████| 28.9k/28.9k [00:00<00:00, 361kB/s]

  0%|          | 0.00/1.65M [00:00<?, ?B/s]
  4%|▍         | 65.5k/1.65M [00:00<00:03, 403kB/s]
 10%|▉         | 164k/1.65M [00:00<00:02, 523kB/s] 
 26%|██▌       | 426k/1.65M [00:00<00:01, 1.03MB/s]
 54%|█████▎    | 885k/1.65M [00:00<00:00, 1.74MB/s]
100%|██████████| 1.65M/1.65M [00:00<00:00, 2.26MB/s]

  0%|          | 0.00/4.54k [00:00<?, ?B/s]
100%|██████████| 4.54k/4.54k [00:00<00:00, 11.0MB/s]
┏━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓
┃   ┃ Name                ┃ Type                ┃ Params ┃ Mode  ┃ FLOPs ┃
┡━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩
│ 0 │ ood_criterion       │ MaxSoftmaxCriterion │      0 │ train │     0 │
│ 1 │ model               │ StochasticModel     │ 88.9 K │ train │     0 │
│ 2 │ loss                │ ELBOLoss            │ 88.9 K │ train │     0 │
│ 3 │ format_batch_fn     │ Identity            │      0 │ train │     0 │
│ 4 │ val_cls_metrics     │ MetricCollection    │      0 │ train │     0 │
│ 5 │ test_cls_metrics    │ MetricCollection    │      0 │ train │     0 │
│ 6 │ test_id_entropy     │ Entropy             │      0 │ train │     0 │
│ 7 │ test_id_ens_metrics │ MetricCollection    │      0 │ train │     0 │
│ 8 │ mixup               │ Identity            │      0 │ train │     0 │
└───┴─────────────────────┴─────────────────────┴────────┴───────┴───────┘
Trainable params: 88.9 K                                                        
Non-trainable params: 0                                                         
Total params: 88.9 K                                                            
Total estimated model params size (MB): 0                                       
Modules in train mode: 68                                                       
Modules in eval mode: 0                                                         
Total FLOPs: 0                                                                  
┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Test metric  ┃      Classification       ┃
┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│     Acc      │          93.490%          │
│    Brier     │          0.09718          │
│   Entropy    │          0.23042          │
│     NLL      │          0.21086          │
└──────────────┴───────────────────────────┘
┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Test metric  ┃        Calibration        ┃
┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│     ECE      │          1.332%           │
│     MCE      │          18.735%          │
│    SmECE     │          1.566%           │
│     aECE     │          1.139%           │
└──────────────┴───────────────────────────┘
┏━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Test metric  ┃ Selective Classification  ┃
┡━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
```

--------------------------------

### Setup Image Downloading and Visualization

Source: https://torch-uncertainty.github.io/auto_tutorials/Data_Augmentation/tutorial_corruption.html

Initializes the environment by downloading sample images and defining helper functions to apply and visualize corruption transforms.

```python
from pathlib import Path
from urllib import request

import matplotlib.pyplot as plt
from PIL import Image
from torchvision.transforms import CenterCrop, Compose, Resize, ToTensor

# Wikipedia URLs (unusable from time to times)
# urls = [
#     "https://upload.wikimedia.org/wikipedia/commons/d/d9/Carduelis_tristis_-Michigan%2C_USA_-male-8.jpg",
#     "https://upload.wikimedia.org/wikipedia/commons/5/5d/Border_Collie_Blanca_y_Negra_Hembra_%28Belen%2C_Border_Collie_Los_Baganes%29.png",
#     "https://upload.wikimedia.org/wikipedia/commons/f/f8/Birmakatze_Seal-Point.jpg",
#     "https://upload.wikimedia.org/wikipedia/commons/a/a9/Garranos_fight.jpg",
#     "https://upload.wikimedia.org/wikipedia/commons/8/8b/Cottontail_Rabbit.jpg",
# ]

# HF URLs (more reliable)
urls = [
    "https://huggingface.co/torch-uncertainty/tutorial-images/resolve/main/Carduelis_tristis_-Michigan%2C_USA_-male-8.jpg",
    "https://huggingface.co/torch-uncertainty/tutorial-images/resolve/main/Border_Collie_Blanca_y_Negra_Hembra_(Belen%2C_Border_Collie_Los_Baganes).png",
    "https://huggingface.co/torch-uncertainty/tutorial-images/resolve/main/Birmakatze_Seal-Point.jpg",
    "https://huggingface.co/torch-uncertainty/tutorial-images/resolve/main/Garranos_fight.jpg",
    "https://huggingface.co/torch-uncertainty/tutorial-images/resolve/main/Cottontail_Rabbit.jpg",
]


def download_img(url, i):
    # Create a request with proper headers to avoid 403 Forbidden error
    if not url.startswith(("http:", "https:")):
        raise ValueError("URL must start with 'http:' or 'https:'")

    req = request.Request(  # noqa: S310
        url,
        headers={
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
        },
    )
    filename = Path(f"tmp_{i}.png")
    with request.urlopen(req) as response, filename.open("wb") as f:  # noqa: S310
        f.write(response.read())
    return Image.open(filename).convert("RGB")


images_ds = [download_img(url, i) for i, url in enumerate(urls)]


def get_images(main_corruption, index: int = 0):
    """Create an image showing the 6 levels of corruption of a given transform."""
    images = []
    for severity in range(6):
        transforms = Compose(
            [
                Resize(256, antialias=True),
                CenterCrop(256),
                ToTensor(),
                main_corruption(severity),
                CenterCrop(224),
            ]
        )
        images.append(transforms(images_ds[index]).permute(1, 2, 0).numpy())
    return images


def show_images(transforms) -> None:
    """Show the effect of all given transforms."""
    num_corruptions = len(transforms)
    _, ax = plt.subplots(num_corruptions, 6, figsize=(10, int(1.5 * num_corruptions)))
    for i, transform in enumerate(transforms):
        images = get_images(transform, index=i)
        ax[i][0].text(
            -0.1,
            0.5,
            transform.__name__,
            transform=ax[i][0].transAxes,
            rotation="vertical",
            horizontalalignment="right",
            verticalalignment="center",
            fontsize=12,
        )
        for j in range(6):
            ax[i][j].imshow(images[j])
            if i == 0 and j == 0:
                ax[i][j].set_title("Original")
            elif i == 0:
                ax[i][j].set_title(f"Severity {j}")
            ax[i][j].axis("off")
    plt.show()
```

--------------------------------

### Validation Start Hook

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/routines/segmentation.html

Actions performed at the start of validation.

```APIDOC
## Validation Start Hook

### Description
Executes actions at the beginning of the validation phase, such as updating the model wrapper if it requires epoch-based updates.

### Method
`on_validation_start`

### Parameters
None

### Response
Updates the model wrapper if `self.needs_epoch_update` is true and not in sanity checking mode.

```

--------------------------------

### Implement Data Preparation and Setup

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datamodules/segmentation/camvid.html

Handles dataset downloading and splitting into train, validation, and test sets based on the provided stage.

```python
    def prepare_data(self) -> None:  # coverage: ignore
        self.dataset(root=self.root, download=True)

    def setup(self, stage: str | None = None) -> None:
        if stage == "fit" or stage is None:
            self.train = self.dataset(
                root=self.root,
                split="train",
                group_classes=self.group_classes,
                download=False,
                transforms=self.train_transform,
            )
            self.val = self.dataset(
                root=self.root,
                split="val",
                group_classes=self.group_classes,
                download=False,
                transforms=self.test_transform,
            )
        if stage == "test" or stage is None:
            self.test = self.dataset(
                root=self.root,
                split="test",
                group_classes=self.group_classes,
                download=False,
                transforms=self.test_transform,
            )

        if stage not in ["fit", "test", None]:
            raise ValueError(f"Stage {stage} is not supported.")
```

--------------------------------

### Training Start Hook

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/routines/segmentation.html

Actions performed at the start of training.

```APIDOC
## Training Start Hook

### Description
Executes actions when training begins, such as logging hyperparameters and checking for a defined loss function.

### Method
`on_train_start`

### Parameters
None

### Response
Raises `ValueError` if the loss function is not specified.

```

--------------------------------

### Install Package in Editable Mode

Source: https://torch-uncertainty.github.io/_sources/installation.rst.txt

Install the package from source in editable mode for development purposes.

```bash
pip install -e .
```

--------------------------------

### Prepare and Setup DataModule for Testing

Source: https://torch-uncertainty.github.io/_downloads/a25e4abf9a4ffd96297cfb7a66bc8ea8/tutorial_distribution_shift.ipynb

Prepares the data and sets up the DataModule for testing. This step is typically handled by the Trainer but is explicitly called here to access the dataset before training.

```python
# These lines are usually not necessary (they are called by the Trainer),
# but we want to get access to the dataset before training
datamodule.prepare_data()
datamodule.setup("test")

# Let's check the CIFAR-10C, it should contain (15+4)*10000 images for the selected severity level.
print(datamodule.shift)
```

--------------------------------

### Build Documentation Locally

Source: https://torch-uncertainty.github.io/_sources/contributing.rst.txt

Navigate to the ./docs directory and run this command to build the HTML documentation locally. Use 'html-noplot' to skip tutorial execution.

```bash
make html
```

--------------------------------

### Mixup Augmentation Setup

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/routines/classification.html

Details the setup and application of Mixup augmentation, including kernel warping variants.

```APIDOC
## Mixup Augmentation Setup

### Description
Sets up the optional Mixup augmentation based on provided parameters. It determines if kernel warping is used and whether it should be applied on embeddings.

### Method
`_init_mixup(self, mixup_params: dict | None) -> Callable`

### Endpoint
None

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **mixup_params** (dict | None) - Optional - The detailed parameters of the mixup augmentation. None if unused.

### Request Example
```json
{
  "mixup_params": {
    "mixtype": "kernel_warping",
    "kw_on_embeddings": true
  }
}
```

### Response
- **Callable** - A callable function that applies the mixup augmentation.

#### Success Response (200)
None

#### Response Example
None
```

```APIDOC
## Apply Mixup Augmentation

### Description
Applies the configured Mixup augmentation to a batch of images and their corresponding targets. Handles different Mixup types, including kernel warping on embeddings or images.

### Method
`_apply_mixup(self, batch: tuple[Tensor, Tensor]) -> tuple[Tensor, Tensor]`

### Endpoint
None

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **batch** (tuple[Tensor, Tensor]) - The images and the corresponding targets.

### Request Example
```json
{
  "batch": [
    "<image_tensor>",
    "<target_tensor>"
  ]
}
```

### Response
- **tuple[Tensor, Tensor]** - The images and the corresponding targets transformed with mixup.

#### Success Response (200)
None

#### Response Example
None
```

--------------------------------

### Initialize Trainer and Data

Source: https://torch-uncertainty.github.io/_downloads/911fedea902d5fd5148110aa3641e71c/tutorial_evidential_classification.ipynb

Sets up the trainer, MNIST datamodule, and LeNet model architecture.

```python
trainer = TUTrainer(accelerator="gpu", devices=1, max_epochs=MAX_EPOCHS, enable_progress_bar=False)

# datamodule
root = Path() / "data"
datamodule = MNISTDataModule(root=root, batch_size=BATCH_SIZE, num_workers=8)

model = lenet(
    in_channels=datamodule.num_channels,
    num_classes=datamodule.num_classes,
)
```

--------------------------------

### Initialize and Run TUTrainer

Source: https://torch-uncertainty.github.io/_downloads/a25e4abf9a4ffd96297cfb7a66bc8ea8/tutorial_distribution_shift.ipynb

Set up the TUTrainer with training parameters like max epochs, precision, and device. Then, use `trainer.fit()` to begin training the model with the defined routine and datamodule.

```python
# Initialize the TUTrainer with a maximum of 10 epochs and the specified device
trainer = TUTrainer(
    max_epochs=10, precision="16-mixed", accelerator="cuda", devices=1, enable_progress_bar=False
)

# Begin training the model using the CIFAR-10 DataModule
trainer.fit(routine, datamodule=datamodule)
```

--------------------------------

### Training Progress Bar Example

Source: https://torch-uncertainty.github.io/auto_tutorials/Classification/tutorial_classification.html

This is an example of a progress bar typically shown during model training, indicating the current status and data transfer rate.

```text
  0%|          | 0.00/9.91M [00:00<?, ?B/s]
  1%|          | 65.5k/9.91M [00:00<00:23, 411kB/s]
  3%|▎         | 262k/9.91M [00:00<00:10, 887kB/s]
  5%|▌         | 524k/9.91M [00:00<00:06, 1.38MB/s]
 10%|▉         | 983k/9.91M [00:00<00:03, 2.36MB/s]
 19%|█▉        | 1.87M/9.91M [00:00<00:01, 4.35MB/s]
 24%|██▍       | 2.36M/9.91M [00:00<00:01, 4.25MB/s]
 29%|██▉       | 2.92M/9.91M [00:00<00:01, 4.60MB/s]
 36%|███▌      | 3.54M/9.91M [00:00<00:01, 5.07MB/s]
 41%|████▏     | 4.10M/9.91M [00:01<00:01, 4.78MB/s]
 47%|████▋     | 4.65M/9.91M [00:01<00:01, 4.86MB/s]
 54%|█████▍    | 5.34M/9.91M [00:01<00:00, 5.11MB/s]
 60%|██████    | 5.96M/9.91M [00:01<00:00, 5.37MB/s]
 66%|██████▌   | 6.55M/9.91M [00:01<00:00, 5.51MB/s]
 72%|███████▏  | 7.14M/9.91M [00:01<00:00, 5.31MB/s]
 79%|███████▊  | 7.80M/9.91M [00:01<00:00, 5.60MB/s]
 85%|████████▌ | 8.45M/9.91M [00:01<00:00, 5.87MB/s]
 92%|█████████▏| 9.08M/9.91M [00:01<00:00, 5.63MB/s]
 98%|█████████▊| 9.73M/9.91M [00:02<00:00, 5.58MB/s]
100%|██████████| 9.91M/9.91M [00:02<00:00, 4.71MB/s]
```

--------------------------------

### Example Output of Model Testing

Source: https://torch-uncertainty.github.io/auto_tutorials/Bayesian_Methods/tutorial_bayesian.html

Provides an example of the expected output format for ground truth, logit shapes, predicted digits, and the standard deviation of scores over posterior samples.

```text
Ground truth:  7 2 1 0
Output logit shape (Num predictions x Batch) x Classes:  torch.Size([64, 10])
Predicted digits:  7 2 1 0
Std. dev. of the scores over the posterior samples 0.000 0.002 0.000 0.012

```

--------------------------------

### Initialize Trainer and DataModule

Source: https://torch-uncertainty.github.io/_downloads/1d73e10cebb64090359837d6682d2001/tutorial_classification.ipynb

Sets up the hardware-aware trainer and the MNIST data provider.

```python
trainer = TUTrainer(accelerator="gpu", devices=1, max_epochs=2, enable_progress_bar=False)

# datamodule providing the dataloaders to the trainer
root = Path("data")
datamodule = MNISTDataModule(root=root, batch_size=128)
```

--------------------------------

### UCIRegressionDataModule Methods

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datamodules/uci_regression.html

Details the prepare_data and setup methods for the UCIRegressionDataModule.

```APIDOC
## UCIRegressionDataModule Methods

### prepare_data

#### Description
Downloads the dataset if it does not exist.

#### Method
prepare_data

### setup

#### Description
Splits the dataset into training, validation, and testing sets.

#### Method
setup

#### Parameters
- **stage** (str | None) - Optional - The stage for which to set up the data (e.g., 'fit', 'test'). Defaults to ``None``.

### Request Example (prepare_data)
```python
datamodule = UCIRegressionDataModule(...) 
datamodule.prepare_data()
```

### Request Example (setup)
```python
datamodule = UCIRegressionDataModule(...) 
datamodule.setup()
```

### Response Example (setup)
After setup, the datamodule will have attributes `train`, `val`, and `test` which are PyTorch Datasets or Subset objects.
```

--------------------------------

### Setup CIFAR-100 Datamodule and Dataloaders

Source: https://torch-uncertainty.github.io/auto_tutorials/Post_Hoc_Methods/tutorial_temperature.html

Initializes the CIFAR-100 datamodule, prepares data, and sets up the test dataloader. Handles cases where out-of-distribution data evaluation is enabled.

```python
dm = CIFAR100DataModule(root="./data", eval_ood=False, batch_size=32)
dm.prepare_data()
dm.setup("test")

```

--------------------------------

### Configuration File Example (config.yaml)

Source: https://torch-uncertainty.github.io/cli_guide.html

This YAML file defines hyperparameters for training, including evaluation settings, trainer configurations, model details, data parameters, optimizer, and learning rate scheduler. It's used to customize the CLI execution.

```yaml
# config.yaml
eval_after_fit: true
trainer:
    max_epochs: 75
    accelerators: gpu
    devices: 1
routine:
    model:
      class_path: torch_uncertainty.models.classification.resnet
      init_args:
      in_channels: 3
      num_classes: 10
      arch: 18
      style: cifar
    num_classes: 10
    loss: CrossEntropyLoss
data:
    root: ./data
    batch_size: 128
optimizer:
  class_path: torch.optim.SGD
  init_args:
    lr: 0.05
    momentum: 0.9
    weight_decay: 5e-4
lr_scheduler:
  class_path: torch.optim.lr_scheduler.MultiStepLR
  init_args:
    milestones:
      - 25
      - 50
    gamma: 0.1

```

--------------------------------

### Setup Training, Validation, and Test Datasets

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datamodules/classification/tiny_imagenet.html

Configures the training, validation, and test datasets based on the specified stage ('fit' or 'test'). It handles train-validation splits and prepares OOD and shifted datasets if enabled.

```python
if stage == "fit" or stage is None:
            full = self.dataset(
                self.root,
                split="train",
                transform=self.train_transform,
            )
            if self.val_split:
                self.train, self.val = create_train_val_split(
                    full,
                    self.val_split,
                    self.test_transform,
                )
            else:
                self.train = full
                self.val = self.dataset(
                    self.root,
                    split="val",
                    transform=self.test_transform,
                )
        if stage == "test" or stage is None:
            self.test = self.dataset(
                self.root,
                split="val",
                transform=self.test_transform,
            )
        if stage not in ["fit", "test", None]:
            raise ValueError(f"Stage {stage} is not supported.")

        if self.eval_ood:
            self.ood = self.ood_dataset(
                self.root,
                split="test",
                transform=self.test_transform,
            )

        if self.eval_shift:
            self.shift = self.shift_dataset(
                self.root,
                download=False,
                shift_severity=self.shift_severity,
                transform=self.test_transform,
            )
```

--------------------------------

### Install Torch-Uncertainty with Dev Packages

Source: https://torch-uncertainty.github.io/contributing.html

Install torch-uncertainty with development packages. Use `--extra gpu` for GPU systems or `--extra cpu` if no GPUs are available. Failure to include the extra flag may result in PyTorch without CUDA, causing issues.

```bash
uv sync --extra gpu
```

```bash
uv sync --extra cpu
```

--------------------------------

### Configure Optimizer and Trainer

Source: https://torch-uncertainty.github.io/_downloads/ab91300a901b0b50ec4f140c171dae8d/tutorial_muad_seg.ipynb

Sets up the optimizer, learning rate scheduler, segmentation routine, and trainer instance.

```python
# We build the optimizer
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)

# Learning rate decay scheduler
lr_updater = lr_scheduler.StepLR(optimizer, step_size=LR_DECAY_EPOCHS, gamma=LR_DECAY)

# Segmentation Routine
seg_routine = SegmentationRoutine(
    model=model,
    num_classes=datamodule.num_classes,
    loss=torch.nn.CrossEntropyLoss(weight=class_weights),
    optim_recipe={"optimizer": optimizer, "lr_scheduler": lr_updater},
)

trainer = TUTrainer(
    accelerator="gpu",
    devices=1,
    max_epochs=NB_EPOCHS,
    enable_progress_bar=False,
    precision="16-mixed",
)
```

--------------------------------

### Update Torch Uncertainty

Source: https://torch-uncertainty.github.io/_sources/installation.rst.txt

Command to upgrade the existing package installation.

```bash
pip install -U torch-uncertainty
```

--------------------------------

### Configure Training Routine

Source: https://torch-uncertainty.github.io/_downloads/c600403fd848d50cfcc5f80a21319f6b/tutorial_muad_deep_en.ipynb

Sets up the optimizer, learning rate scheduler, segmentation routine, and trainer.

```python
# We build the optimizer
optimizer = optim.Adam(
    ensemble.parameters(), lr=LEARNING_RATE * NUM_ESTIMATORS, weight_decay=WEIGHT_DECAY
)

# Learning rate decay scheduler
lr_updater = lr_scheduler.StepLR(optimizer, LR_DECAY_EPOCHS, LR_DECAY)

ens_routine = SegmentationRoutine(
    model=ensemble,
    num_classes=datamodule.num_classes,
    loss=torch.nn.CrossEntropyLoss(weight=class_weights),
    format_batch_fn=RepeatTarget(NUM_ESTIMATORS),  # Repeat the target 4 times for the ensemble
    optim_recipe={"optimizer": optimizer, "lr_scheduler": lr_updater},
)

trainer = TUTrainer(
    accelerator="gpu",
    devices=1,
    max_epochs=NB_EPOCHS,
    enable_progress_bar=False,
    precision="16-mixed",
)
```

--------------------------------

### GET /dataset/camvid

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datasets/segmentation/camvid.html

Retrieves a specific sample from the CamVid dataset by index.

```APIDOC
## GET /dataset/camvid

### Description
Retrieves the image and target mask at the specified index from the dataset.

### Method
GET

### Parameters
#### Query Parameters
- **index** (int) - Required - The sample index to retrieve.

### Response
#### Success Response (200)
- **image** (tv_tensors.Image) - The loaded image.
- **target** (tv_tensors.Mask) - The encoded target mask.
```

--------------------------------

### Lifecycle Hooks

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/routines/classification.html

Methods for handling training, validation, and testing start events.

```python
    def on_train_start(self) -> None:  # coverage: ignore
        """Put the hyperparameters in tensorboard."""
        if self.loss is None:
            raise ValueError(
                "To train a model, you must specify the `loss` argument in the routine. Got None."
            )
        if self.logger is not None:
            self.logger.log_hyperparams(self.hparams)
```

```python
    def on_validation_start(self) -> None:
        """Prepare the validation step.

        Update the model's wrapper and the batchnorms if needed.
        """
        if self.needs_epoch_update and not self.trainer.sanity_checking:
            self.model.update_wrapper(self.current_epoch)
            if hasattr(self.model, "need_bn_update"):
                self.model.bn_update(self.trainer.train_dataloader, device=self.device)
```

```python
    def on_test_start(self) -> None:
        """Prepare the test step.

        Setup the post-processing dataset and fit the post-processing method if needed, prepares
        the storage lists for logit plotting and update the batchnorms if needed.
        """
        if self.post_processing is not None:
            with torch.inference_mode(False):
                self.post_processing.fit(self.trainer.datamodule.postprocess_dataloader())

        if self.eval_ood and self.log_plots and isinstance(self.logger, Logger):
```

--------------------------------

### Get Dataset Length

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datasets/muad.html

Returns the total number of samples available in the dataset.

```APIDOC
## Get Dataset Length

### Description
Returns the total number of samples in the dataset.

### Method
```python
__len__(self) -> int
```

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Request Example
```python
num_samples = len(dataset)
```

### Response
#### Success Response (200)
- **int** - The total number of samples in the dataset.

#### Response Example
```json
{
  "num_samples": 1000
}
```
```

--------------------------------

### CLI Subcommand Help Output

Source: https://torch-uncertainty.github.io/cli_guide.html

Example output showing the available subcommands for the CLI tool, including 'fit', 'validate', 'test', and 'predict'. Each subcommand can be further explored with '--help'.

```bash
subcommands:
For more details of each subcommand, add it as an argument followed by --help.

Available subcommands:
    fit                 Runs the full optimization routine.
    validate            Perform one evaluation epoch over the validation set.
    test                Perform one evaluation epoch over the test set.
    predict             Run evaluation on your data.


```

--------------------------------

### GET __getitem__

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datasets/regression/uci_regression.html

Retrieves a sample and its corresponding target from the dataset at a specified index.

```APIDOC
## GET __getitem__

### Description
Retrieves a sample and target for a given index. If the dataset is 'energy-prediction', it returns a slice of the data; otherwise, it returns the item at the index with optional transformations applied.

### Method
GET

### Parameters
#### Path Parameters
- **index** (int) - Required - The index of the data point to retrieve.

### Response
#### Success Response (200)
- **data** (torch.Tensor) - The retrieved data sample.
- **target** (torch.Tensor) - The corresponding target value.
```

--------------------------------

### Initialize Trainer, Data, and Model

Source: https://torch-uncertainty.github.io/_downloads/d13ff3c657c4f5e710ba744815d554a5/tutorial_mc_batch_norm.ipynb

Sets up the PyTorch Lightning trainer, MNIST data module, and LeNet architecture with batch normalization.

```python
trainer = TUTrainer(accelerator="gpu", devices=1, max_epochs=2, enable_progress_bar=False)

# datamodule
root = Path("data")
datamodule = MNISTDataModule(root, batch_size=128)


model = lenet(
    in_channels=datamodule.num_channels,
    num_classes=datamodule.num_classes,
    norm=nn.BatchNorm2d,
)
```

--------------------------------

### Get Dataset Length

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/datasets/classification/imagenet/tiny_imagenet.html

Returns the total number of samples available in the dataset.

```python
    def __len__(self) -> int:
        """The number of samples in the dataset."""
        return self.samples_num

```

--------------------------------

### Handle Validation Start

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/routines/pixel_regression.html

Updates the model wrapper and batch normalization layers if required.

```python
    def on_validation_start(self) -> None:
        """Prepare the validation step.

        Update the model's wrapper and the batchnorms if needed.
        """
        if self.needs_epoch_update and not self.trainer.sanity_checking:
            self.model.update_wrapper(self.current_epoch)
            if hasattr(self.model, "need_bn_update"):
                self.model.bn_update(self.trainer.train_dataloader, device=self.device)
```

--------------------------------

### Initialize Trainer, Datamodule, and Model

Source: https://torch-uncertainty.github.io/auto_tutorials/Classification/tutorial_evidential_classification.html

Sets up the training environment by initializing the TUTrainer, MNISTDataModule, and the LeNet model. The trainer is configured for GPU acceleration.

```python
trainer = TUTrainer(accelerator="gpu", devices=1, max_epochs=MAX_EPOCHS, enable_progress_bar=False)

# datamodule
root = Path() / "data"
datamodule = MNISTDataModule(root=root, batch_size=BATCH_SIZE, num_workers=8)

model = lenet(
    in_channels=datamodule.num_channels,
    num_classes=datamodule.num_classes,
)

```

--------------------------------

### Initialize and use AdaptiveCalibrationError

Source: https://torch-uncertainty.github.io/_modules/torch_uncertainty/metrics/classification/calibration/adaptive_calibration_error.html

Example demonstrating the usage of AdaptiveCalibrationError for a binary classification task.

```python
from torch_uncertainty.metrics.classification.adaptive_calibration_error import (
    AdaptiveCalibrationError,
)

# Binary classification example
predicted_probs = torch.tensor([0.95, 0.85, 0.15, 0.05])
true_labels = torch.tensor([1, 1, 0, 0])

metric = CalibrationError(
    task="binary",
    num_bins=5,
    norm="l1",
)

calibration_error = metric(predicted_probs, true_labels)
print(f"Calibration Error (Binary): {calibration_error}")
# Output : Calibration Error (Binary): 0.1
```
Tokens:
Raw