Shortcuts

Source code for torch_uncertainty.post_processing.calibration.vector_scaler

from typing import Literal

import torch
from torch import nn

from .scaler import Scaler


[docs]class VectorScaler(Scaler): def __init__( self, num_classes: int, model: nn.Module | None = None, init_w: float = 1, init_b: float = 0, lr: float = 0.1, max_iter: int = 200, device: Literal["cpu", "cuda"] | torch.device | None = None, ) -> None: """Vector scaling post-processing for calibrated probabilities. Args: model (nn.Module): Model to calibrate. num_classes (int): Number of classes. init_w (float, optional): Initial value for the weights. Defaults to 1. init_b (float, optional): Initial value for the bias. Defaults to 0. lr (float, optional): Learning rate for the optimizer. Defaults to 0.1. max_iter (int, optional): Maximum number of iterations for the optimizer. Defaults to 100. device (Optional[Literal["cpu", "cuda"]], optional): Device to use for optimization. Defaults to None. Reference: Guo, C., Pleiss, G., Sun, Y., & Weinberger, K. Q. On calibration of modern neural networks. In ICML 2017. """ super().__init__(model=model, lr=lr, max_iter=max_iter, device=device) if not isinstance(num_classes, int): raise TypeError("num_classes must be an integer.") if num_classes <= 0: raise ValueError("The number of classes must be positive.") self.num_classes = num_classes self.set_temperature(init_w, init_b)
[docs] def set_temperature(self, val_w: float, val_b: float) -> None: """Set the temperature to a fixed value. Args: val_w (float): Weight temperature value. val_b (float): Bias temperature value. """ self.temp_w = nn.Parameter( torch.ones(self.num_classes, device=self.device) * val_w, requires_grad=True, ) self.temp_b = nn.Parameter( torch.ones(self.num_classes, device=self.device) * val_b, requires_grad=True, )
def _scale(self, logits: torch.Tensor) -> torch.Tensor: """Scale the predictions with the optimal temperature. Args: logits (torch.Tensor): logits to be scaled. Returns: torch.Tensor: Scaled logits. """ return self.temp_w * logits + self.temp_b @property def temperature(self) -> list: return [self.temp_w, self.temp_b]