Skip to content

Core Abstractions

The core module contains the central Ensemble class and adapters that allow Bensemble to work with any PyTorch model.

Ensemble

bensemble.core.ensemble.Ensemble 

Ensemble(
    members: MemberAdapter,
    combiner: Callable[[Tensor], Tensor] | None = None,
)

Bases: Module

Ensemble of model predictors combined via an aggregation function.

Initializes the Ensemble container.

Parameters:

Name Type Description Default
members MemberAdapter

MemberAdapter instance wrapping ensemble members.

 required
combiner Callable[[Tensor], Tensor] | None

Aggregation function taking (M, batch, ) and returning combined predictions (batch, ). Defaults to averaging across members.

 None
Source code in bensemble/core/ensemble.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
def __init__(
    self,
    members: MemberAdapter,
    combiner: Callable[[torch.Tensor], torch.Tensor] | None = None,
):
    """Initializes the Ensemble container.

    Args:
        members: MemberAdapter instance wrapping ensemble members.
        combiner: Aggregation function taking (M, batch, *) and returning
            combined predictions (batch, *). Defaults to averaging across members.
    """
    super().__init__()
    self.members = members
    self.combiner = combiner or (lambda preds: preds.mean(dim=0))

member_modules property 

member_modules: list[Module]

list[nn.Module]: Underlying member modules.

num_members property 

num_members: int

int: Number of ensemble members.

forward 

forward(x: Tensor) -> Predictions

Computes combined ensemble prediction.

Parameters:

Name Type Description Default
x Tensor

Input tensor of shape (batch_size, *).

 required

Returns:

Name Type Description
Predictions Predictions

Aggregated output prediction tensor.
Source code in bensemble/core/ensemble.py 
40
41
42
43
44
45
46
47
48
49
def forward(self, x: torch.Tensor) -> Predictions:
    """Computes combined ensemble prediction.

    Args:
        x: Input tensor of shape (batch_size, *).

    Returns:
        Predictions: Aggregated output prediction tensor.
    """
    return self.combiner(self.predict_members(x))

from_models classmethod 

from_models(models: list[Module]) -> Ensemble

Creates an explicit ensemble from a list of independent models.

Parameters:

Name Type Description Default
models list[Module]

List of trained PyTorch modules.

 required

Returns:

Name Type Description
Ensemble Ensemble

Configured Ensemble instance.
Source code in bensemble/core/ensemble.py 
61
62
63
64
65
66
67
68
69
70
71
@classmethod
def from_models(cls, models: list[nn.Module]) -> "Ensemble":
    """Creates an explicit ensemble from a list of independent models.

    Args:
        models: List of trained PyTorch modules.

    Returns:
        Ensemble: Configured Ensemble instance.
    """
    return cls(members=ExplicitMembers(models))

from_posterior classmethod 

from_posterior(
    source: PosteriorSource, n_members: int = 10, **kwargs
) -> Ensemble

Creates an explicit ensemble by sampling from an approximated posterior.

Parameters:

Name Type Description Default
source PosteriorSource

Posterior approximation source implementing sample_models.

 required
n_members int

Number of members to sample. Defaults to 10.

 10
**kwargs

Additional keyword arguments passed to sample_models.

 {}

Returns:

Name Type Description
Ensemble Ensemble

Configured Ensemble instance.
Source code in bensemble/core/ensemble.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
@classmethod
def from_posterior(
    cls, source: PosteriorSource, n_members: int = 10, **kwargs
) -> "Ensemble":
    """Creates an explicit ensemble by sampling from an approximated posterior.

    Args:
        source: Posterior approximation source implementing `sample_models`.
        n_members: Number of members to sample. Defaults to 10.
        **kwargs: Additional keyword arguments passed to `sample_models`.

    Returns:
        Ensemble: Configured Ensemble instance.
    """
    models = source.sample_models(n_members, **kwargs)
    return cls(members=ExplicitMembers(models))

from_stochastic classmethod 

from_stochastic(
    model: Module, num_samples: int = 30, mode: str = "auto"
) -> Ensemble

Creates an implicit ensemble from a single stochastic model.

Parameters:

Name Type Description Default
model Module

Stochastic neural network module.

 required
num_samples int

Number of stochastic forward samples. Defaults to 30.

 30
mode str

Stochastic mode ("dropout", "bayesian", or "auto"). Defaults to "auto".

 'auto'

Returns:

Name Type Description
Ensemble Ensemble

Configured Ensemble instance.
Source code in bensemble/core/ensemble.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
@classmethod
def from_stochastic(
    cls, model: nn.Module, num_samples: int = 30, mode: str = "auto"
) -> "Ensemble":
    """Creates an implicit ensemble from a single stochastic model.

    Args:
        model: Stochastic neural network module.
        num_samples: Number of stochastic forward samples. Defaults to 30.
        mode: Stochastic mode ("dropout", "bayesian", or "auto"). Defaults to "auto".

    Returns:
        Ensemble: Configured Ensemble instance.
    """
    return cls(members=StochasticMembers(model, num_samples, mode))

predict_members 

predict_members(x: Tensor) -> MemberPredictions

Computes predictions for each individual member.

Parameters:

Name Type Description Default
x Tensor

Input tensor of shape (batch_size, *).

 required

Returns:

Name Type Description
MemberPredictions MemberPredictions

Per-member predictions of shape (num_members, batch_size, *).
Source code in bensemble/core/ensemble.py 
29
30
31
32
33
34
35
36
37
38
def predict_members(self, x: torch.Tensor) -> MemberPredictions:
    """Computes predictions for each individual member.

    Args:
        x: Input tensor of shape (batch_size, *).

    Returns:
        MemberPredictions: Per-member predictions of shape (num_members, batch_size, *).
    """
    return self.members.predict_all(x)

Member Adapters

bensemble.core.member.MemberAdapter

Bases: Module

Adapts different prediction sources into a uniform (M, batch, *) interface.

bensemble.core.member.ExplicitMembers 

ExplicitMembers(models: list[Module])

Bases: MemberAdapter

Wraps a list of independent nn.Module instances.

Source code in bensemble/core/member.py 
29
30
31
def __init__(self, models: list[nn.Module]):
    super().__init__()
    self.models = nn.ModuleList(models)

bensemble.core.member.StochasticMembers 

StochasticMembers(
    model: Module, num_samples: int = 30, mode: str = "auto"
)

Bases: MemberAdapter

Wraps a single model whose forward pass is stochastic.

Source code in bensemble/core/member.py 
62
63
64
65
66
def __init__(self, model: nn.Module, num_samples: int = 30, mode: str = "auto"):
    super().__init__()
    self.model = model
    self.num_samples = num_samples
    self.mode = mode if mode != "auto" else self._detect_mode()