torch_uncertainty.models.deep_ensembles

torch_uncertainty.models.deep_ensembles(models, num_estimators=None, task='classification', probabilistic=None, reset_model_parameters=False)

Build a Deep Ensembles out of the original models.

Parameters:

• models (list[nn.Module] | nn.Module) – The model to be ensembled.

• num_estimators (int | None) – The number of estimators in the ensemble.

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

• probabilistic (bool) – Whether the regression model is probabilistic.

• reset_model_parameters (bool) – Whether to reset the model parameters when :attr:models is a module or a list of length 1.

Returns:

The ensembled model.

Return type:

_DeepEnsembles

Raises:

• ValueError – If :attr:num_estimators is not specified and :attr:models is a module (or singleton list).

• ValueError – If :attr:num_estimators is less than 2 and :attr:models is a module (or singleton list).

• ValueError – If :attr:num_estimators is defined while :attr:models is a (non-singleton) list.

References

Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In NeurIPS, 2017.