Calibration

bensemble.calibration.scaling

Post-hoc calibration for classifiers.

Both scalers recalibrate class logits against integer class labels; neither has a regression counterpart.

TemperatureScaling

TemperatureScaling(init_temp: float = 1.5)

Bases: Module

Temperature Scaling for classifier calibration.

Divides logits by a single learnable scalar parameter T (temperature). This softens probabilities and calibrates confidence without changing classification accuracy (argmax remains identical).

Initializes the TemperatureScaling module.

Parameters:

Name	Type	Description	Default
`init_temp`	`float`	Initial value for the temperature scalar. Defaults to 1.5.	`1.5`

Source code in bensemble/calibration/scaling.py

def __init__(self, init_temp: float = 1.5):
    """Initializes the TemperatureScaling module.

    Args:
        init_temp: Initial value for the temperature scalar. Defaults to 1.5.
    """
    super().__init__()
    self.temperature = nn.Parameter(torch.ones(1) * init_temp)

fit

fit(logits: Tensor, labels: Tensor, max_iter: int = 50, optimizer: Optimizer | None = None) -> TemperatureScaling

Finds the optimal temperature T using a validation set.

Optimizes the negative log-likelihood.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Unscaled logits from a hold-out validation set of shape (N, num_classes).	required
`labels`	`Tensor`	Ground truth class indices of shape (N,).	required
`max_iter`	`int`	Maximum number of optimizer iterations. Defaults to 50.	`50`
`optimizer`	`Optimizer \| None`	Optimizer over this module's parameters. Defaults to L-BFGS with a strong Wolfe line search. L-BFGS runs up to `max_iter` iterations within a single step, while any other optimizer is stepped `max_iter` times.	`None`

Returns:

Name	Type	Description
`TemperatureScaling`	`TemperatureScaling`	The fitted instance itself.

Source code in bensemble/calibration/scaling.py

def fit(
    self,
    logits: torch.Tensor,
    labels: torch.Tensor,
    max_iter: int = 50,
    optimizer: optim.Optimizer | None = None,
) -> "TemperatureScaling":
    """Finds the optimal temperature T using a validation set.

    Optimizes the negative log-likelihood.

    Args:
        logits: Unscaled logits from a hold-out validation set of shape
            (N, num_classes).
        labels: Ground truth class indices of shape (N,).
        max_iter: Maximum number of optimizer iterations. Defaults to 50.
        optimizer: Optimizer over this module's parameters. Defaults to
            L-BFGS with a strong Wolfe line search. L-BFGS runs up to
            `max_iter` iterations within a single step, while any other
            optimizer is stepped `max_iter` times.

    Returns:
        TemperatureScaling: The fitted instance itself.
    """
    _fit_calibrator(self, [self.temperature], logits, labels, max_iter, optimizer)
    return self

forward

forward(logits: Tensor) -> torch.Tensor

Applies temperature scaling to the input logits.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Raw uncalibrated logits of shape (batch_size, num_classes).	required

Returns:

Type	Description
`Tensor`	torch.Tensor: Scaled logits of shape (batch_size, num_classes).

Source code in bensemble/calibration/scaling.py

def forward(self, logits: torch.Tensor) -> torch.Tensor:
    """Applies temperature scaling to the input logits.

    Args:
        logits: Raw uncalibrated logits of shape (batch_size, num_classes).

    Returns:
        torch.Tensor: Scaled logits of shape (batch_size, num_classes).
    """
    return logits / self.temperature

VectorScaling

VectorScaling(num_classes: int)

Bases: Module

Vector Scaling for multi-class calibration (extension of Platt Scaling).

Applies a per-class affine transformation to uncalibrated logits: calibrated_logits = logits * a + b

Initializes the VectorScaling module.

Parameters:

Name	Type	Description	Default
`num_classes`	`int`	Number of classes in the classification task.	required

Source code in bensemble/calibration/scaling.py

def __init__(self, num_classes: int):
    """Initializes the VectorScaling module.

    Args:
        num_classes: Number of classes in the classification task.
    """
    super().__init__()
    self.a = nn.Parameter(torch.ones(num_classes))
    self.b = nn.Parameter(torch.zeros(num_classes))

fit

fit(logits: Tensor, labels: Tensor, max_iter: int = 50, optimizer: Optimizer | None = None) -> VectorScaling

Finds optimal scaling vectors 'a' and 'b' using a validation set.

Optimizes the negative log-likelihood.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Unscaled logits from a hold-out validation set of shape (N, num_classes).	required
`labels`	`Tensor`	Ground truth class indices of shape (N,).	required
`max_iter`	`int`	Maximum number of optimizer iterations. Defaults to 50.	`50`
`optimizer`	`Optimizer \| None`	Optimizer over this module's parameters. Defaults to L-BFGS with a strong Wolfe line search. L-BFGS runs up to `max_iter` iterations within a single step, while any other optimizer is stepped `max_iter` times.	`None`

Returns:

Name	Type	Description
`VectorScaling`	`VectorScaling`	The fitted instance itself.

Source code in bensemble/calibration/scaling.py

def fit(
    self,
    logits: torch.Tensor,
    labels: torch.Tensor,
    max_iter: int = 50,
    optimizer: optim.Optimizer | None = None,
) -> "VectorScaling":
    """Finds optimal scaling vectors 'a' and 'b' using a validation set.

    Optimizes the negative log-likelihood.

    Args:
        logits: Unscaled logits from a hold-out validation set of shape
            (N, num_classes).
        labels: Ground truth class indices of shape (N,).
        max_iter: Maximum number of optimizer iterations. Defaults to 50.
        optimizer: Optimizer over this module's parameters. Defaults to
            L-BFGS with a strong Wolfe line search. L-BFGS runs up to
            `max_iter` iterations within a single step, while any other
            optimizer is stepped `max_iter` times.

    Returns:
        VectorScaling: The fitted instance itself.
    """
    _fit_calibrator(self, [self.a, self.b], logits, labels, max_iter, optimizer)
    return self

forward

forward(logits: Tensor) -> torch.Tensor

Applies learned affine transformation to the input logits.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Raw uncalibrated logits of shape (batch_size, num_classes).	required

Returns:

Type	Description
`Tensor`	torch.Tensor: Calibrated logits of shape (batch_size, num_classes).

Source code in bensemble/calibration/scaling.py

def forward(self, logits: torch.Tensor) -> torch.Tensor:
    """Applies learned affine transformation to the input logits.

    Args:
        logits: Raw uncalibrated logits of shape (batch_size, num_classes).

    Returns:
        torch.Tensor: Calibrated logits of shape (batch_size, num_classes).
    """
    return logits * self.a + self.b