TemperatureScaler#
- class torch_uncertainty.post_processing.TemperatureScaler(model=None, init_val=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_val (float, 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.
- fit(dataloader, save_logits=False, progress=True)#
Fit the temperature parameters to the calibration data.
- Parameters:
dataloader (DataLoader) – Dataloader with the calibration data. If there is no model, the dataloader should include the confidence score directly and not the logits.
save_logits (bool, optional) – Whether to save the logits and labels in memory. Defaults to
False
.progress (bool, optional) – Whether to show a progress bar. Defaults to
True
.