mc_dropout#

torch_uncertainty.methods.mc_dropout(core_model, num_estimators, last_layer=False, on_batch=True, task='classification', probabilistic=None)[source]#

MC Dropout wrapper for a model.

Parameters:

  • core_model (Module) – Model to wrap.

  • num_estimators (int) – Number of estimators to use during evaluation.

  • last_layer (bool) – Whether to apply dropout to the last layer only. Defaults to False.

  • on_batch (bool) – Increase the batch size to perform MC-Dropout instead of using a Python loop. Increases memory footprint. Defaults to True.

  • task (Literal['classification', 'regression', 'segmentation', 'pixel_regression']) – The model task. Defaults to "classification".

  • probabilistic (bool | None) – Whether the regression model is probabilistic. Defaults to None

Return type:

_MCDropout

Warning

Beware that on_batch==True can raise weird errors if not enough memory is available.