deep_ensembles#

torch_uncertainty.models.deep_ensembles(models, num_estimators=None, task='classification', probabilistic=None, reset_model_parameters=True, store_on_cpu=False, ckpt_paths=None, use_tu_ckpt_format=False)[source]#

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. Defaults to True.

  • store_on_cpu (bool) – Whether to store the models on CPU. Defaults to False. This is useful for large models that do not fit in GPU memory. Only one model will be stored on GPU at a time during forward. The rest will be stored on CPU.

  • ckpt_paths (list[str | Path] | None) – The paths to the checkpoints of the models. If provided, the models will be loaded from the checkpoints. The number of models and the number of checkpoint paths must be the same. If not provided, the models will be used as is. Defaults to None.

  • use_tu_ckpt_format (bool) – Whether the checkpoint is from torch-uncertainty. If True, the checkpoint will be loaded using the torch-uncertainty loading function. If False, the checkpoint will be loaded using the default PyTorch loading function. Note that this option is only used if :attr:ckpt_paths is provided. Defaults to False.

Returns:

The ensembled model.

Return type:

_DeepEnsembles | _RegDeepEnsembles

Raises:

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

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

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

Warning

The store_on_cpu option is not supported for training. It is only supported for inference.

References

[1] Simple and scalable predictive uncertainty estimation using deep ensembles. In NeurIPS, 2017.

On this page

This Page