batch_ensemble#

torch_uncertainty.methods.batch_ensemble(core_model, num_estimators, repeat_training_inputs=False, convert_layers=False)[source]#

BatchEnsemble wrapper for a model.

Parameters:

  • core_model (Module) – Model to wrap.

  • num_estimators (int) – Number of ensemble members.

  • repeat_training_inputs (bool) – Whether to repeat the input batch during training. If True, the input batch is repeated during both training and evaluation. If False, the input batch is repeated only during evaluation. Defaults to False.

  • convert_layers (bool) – Whether to convert the model’s layers to BatchEnsemble layers. If True, the wrapper will convert all nn.Linear and nn.Conv2d layers to their BatchEnsemble counterparts. Defaults to False.

Returns:

BatchEnsemble wrapper for the core_model

Return type:

BatchEnsemble