MatrixScaler¶
- class torch_uncertainty.post_processing.MatrixScaler(model, num_classes, init_w=1, init_b=0, lr=0.1, max_iter=200, device=None)[source]¶
Matrix scaling post-processing for calibrated probabilities.
- Parameters:
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.
- fit(calibration_set, save_logits=False, progress=True)¶
Fit the temperature parameters to the calibration data.
- Parameters:
calibration_set (Dataset) – Calibration dataset.
save_logits (bool, optional) – Whether to save the logits and labels. Defaults to False.
progress (bool, optional) – Whether to show a progress bar. Defaults to True.
- Returns:
Calibrated scaler.
- Return type:
Scaler