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=None)

Bases: Module

A collection of predictors whose outputs are combined. Ensemble doesn't care where its members came from.

Source code in bensemble/core/ensemble.py 
13
14
15
16
def __init__(self, members: MemberAdapter, combiner=None):
    super().__init__()
    self.members = members
    self.combiner = combiner or (lambda preds: preds.mean(dim=0))

member_modules property 

member_modules: list[Module]

Access underlying nn.Modules (for KL, parameters, etc.).

forward 

forward(x: Tensor) -> Predictions

Combined prediction.

Source code in bensemble/core/ensemble.py 
22
23
24
def forward(self, x: torch.Tensor) -> Predictions:
    """Combined prediction."""
    return self.combiner(self.predict_members(x))

from_models classmethod 

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

Explicit ensemble from a list of independent models.

Source code in bensemble/core/ensemble.py 
35
36
37
38
@classmethod
def from_models(cls, models: list[nn.Module]) -> "Ensemble":
    """Explicit ensemble from a list of independent models."""
    return cls(members=ExplicitMembers(models))

from_posterior classmethod 

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

Sample an explicit ensemble from a fitted posterior approximation.

Source code in bensemble/core/ensemble.py 
49
50
51
52
53
54
55
56
57
@classmethod
def from_posterior(
    cls, source: PosteriorSource, n_members: int = 10, **kwargs
) -> "Ensemble":
    """
    Sample an explicit ensemble from a fitted posterior approximation.
    """
    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

Implicit ensemble from a model with stochastic forward passes.

Source code in bensemble/core/ensemble.py 
40
41
42
43
44
45
46
47
@classmethod
def from_stochastic(
    cls, model: nn.Module, num_samples: int = 30, mode: str = "auto"
) -> "Ensemble":
    """
    Implicit ensemble from a model with stochastic forward passes.
    """
    return cls(members=StochasticMembers(model, num_samples, mode))

predict_members 

predict_members(x: Tensor) -> MemberPredictions

(M, batch, *) — raw per-member outputs.

Source code in bensemble/core/ensemble.py 
18
19
20
def predict_members(self, x: torch.Tensor) -> MemberPredictions:
    """(M, batch, *) — raw per-member outputs."""
    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 
26
27
28
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 
53
54
55
56
57
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()