Source code for torch_uncertainty.post_processing.conformal.aps

from typing import Literal

import torch
import torch.nn.functional as F
from torch import Tensor, nn
from torch.utils.data import DataLoader

from .abstract import Conformal




class ConformalClsAPS(Conformal):
    def __init__(
        self,
        alpha: float,
        model: nn.Module | None = None,
        randomized: bool = True,
        ts_init_val: float = 1,
        ts_lr: float = 0.1,
        ts_max_iter: int = 100,
        enable_ts: bool = True,
        device: Literal["cpu", "cuda"] | torch.device | None = None,
    ) -> None:
        r"""Conformal classification with Adaptive Prediction Sets (APS;
        Romano, Sesia & Candès, NeurIPS 2020).

        Uses as non-conformity score the cumulative probability mass needed to reach
        the true class once the predictions are sorted by decreasing probability. For
        a sample with predicted probabilities :math:`\hat{\mathbf{p}}` whose true class
        ranks at position :math:`k` after sorting,

        .. math::
            s(\mathbf{x}, y) = \sum_{i=1}^{k} \hat{p}_{(i)} - U \cdot \hat{p}_{(k)},

        where :math:`U \sim \mathrm{Uniform}(0, 1)` smooths the cumulative score when
        :attr:`randomized` is ``True``. The calibrated quantile :math:`\hat{q}` defines
        the test-time prediction set

        .. math::
            \mathcal{C}(\mathbf{x}) = \{ c : s(\mathbf{x}, c) \leq \hat{q} \},

        which adapts in size to the difficulty of each example: easy points get tight
        sets, ambiguous points get larger ones.

        Args:
            alpha: Target mis-coverage level :math:`\alpha \in (0, 1)`.
            model: Trained classification model. Defaults to ``None``.
            randomized: Whether to use randomised tie-breaking in APS. Defaults to ``True``.
            ts_init_val: Initial value for the temperature. Defaults to ``1.0``.
            ts_lr: Learning rate for the temperature scaling optimizer. Defaults to ``0.1``.
            ts_max_iter: Maximum number of iterations for the temperature scaling optimizer.
                Defaults to ``100``.
            enable_ts: Whether to apply temperature scaling to the logits before
                computing the conformal scores. Defaults to ``True``.
            device: Device to use. Defaults to ``None``.

        Warning:
            This implementation only works in the multiclass setting. Raise an issue
            if binary support is needed.

        Reference:
            - `Romano, Y., Sesia, M., & Candès, E. (2020). Classification with Valid
              and Adaptive Coverage. NeurIPS 2020 <https://arxiv.org/abs/2006.02544>`_.

        Code inspired by TorchCP.
        """
        super().__init__(
            alpha=alpha,
            model=model,
            ts_init_val=ts_init_val,
            ts_lr=ts_lr,
            ts_max_iter=ts_max_iter,
            enable_ts=enable_ts,
            device=device,
        )
        self.randomized = randomized



    def model_forward(self, inputs: Tensor) -> Tensor:
        """Apply the model and return the scores."""
        assert self.model is not None
        self.model.eval()
        return self.model(inputs.to(self.device)).softmax(-1)


    def _sort_sum(self, probs: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        """Sort probabilities and compute cumulative sums."""
        ordered, indices = torch.sort(probs, dim=-1, descending=True)
        cumsum = torch.cumsum(ordered, dim=-1)
        return indices, ordered, cumsum

    def _calculate_all_labels(self, probs: Tensor) -> Tensor:
        """Calculate APS scores for all labels."""
        indices, ordered, cumsum = self._sort_sum(probs)
        if self.randomized:
            noise = torch.rand(probs.shape, device=probs.device)
        else:
            noise = torch.zeros_like(probs)

        ordered_scores = cumsum - ordered * noise
        _, sorted_indices = torch.sort(indices, descending=False, dim=-1)
        return ordered_scores.gather(dim=-1, index=sorted_indices)

    def _calculate_single_label(self, probs: Tensor, label: Tensor) -> Tensor:
        """Calculate APS score for a single label."""
        indices, ordered, cumsum = self._sort_sum(probs)
        if self.randomized:
            noise = torch.rand(indices.shape[0], device=probs.device)
        else:
            noise = torch.zeros(indices.shape[0], device=probs.device)

        idx = torch.where(indices == label.view(-1, 1))
        return cumsum[idx] - noise * ordered[idx]



    @torch.no_grad()
    def fit(self, dataloader: DataLoader) -> None:
        """Calibrate the APS threshold q_hat on a calibration set."""
        assert self.model is not None
        if self.enable_ts:
            self.model.fit(dataloader=dataloader)

        aps_scores = []
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(self.device), labels.to(self.device)
            probs = self.model_forward(inputs)
            scores = self._calculate_single_label(probs, labels)
            aps_scores.append(scores)

        self.q_hat = torch.quantile(torch.cat(aps_scores), 1 - self.alpha)




    @torch.no_grad()
    def conformal(self, inputs: Tensor) -> Tensor:
        """Compute the prediction set for each input.

        Returns:
            A probability vector for each sample. Its support is the conformal
            prediction set, with equal probability assigned to every included class.
            If thresholding produces an empty prediction set, the class with the
            lowest non-conformity score is included before normalization. Consequently,
            every returned row is finite and sums to one.
        """
        probs = self.model_forward(inputs)
        scores = self._calculate_all_labels(probs)
        pred_set = scores <= self.quantile

        empty = ~pred_set.any(dim=-1)
        fallback = F.one_hot(
            scores.argmin(dim=-1),
            num_classes=scores.shape[-1],
        ).bool()
        pred_set = pred_set | (empty.unsqueeze(-1) & fallback)

        pred_set = pred_set.to(dtype=probs.dtype)
        return pred_set / pred_set.sum(dim=-1, keepdim=True)