TemperatureScaler#

class torch_uncertainty.post_processing.TemperatureScaler(model=None, init_temperature=1, lr=0.1, max_iter=100, eps=1e-08, device=None)[source]#

Temperature scaling post-processing for calibrated probabilities.

Parameters:

  • model (nn.Module) – Model to calibrate.

  • init_temperature (float | Tensor, optional) – Initial value for the temperature. Defaults to 1.

  • 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.

  • eps (float) – Small value for stability. Defaults to 1e-8.

  • device (Optional[Literal["cpu", "cuda"]], optional) – Device to use for optimization. Defaults to None.

References

[1] On calibration of modern neural networks. In ICML 2017.

Warning

If the model is binary, we will by default apply the sigmoid before transposing the prediction to the corresponding 2-class logits.

Note

The Scaler will log an error if the temperature after fitting is negative.

set_temperature(val)[source]#

Set the temperature to a fixed value.

Parameters:

val (float | Tensor) – Temperature value.