Ensemble Fit Manager#

`fit_manager` #

CLASS	DESCRIPTION
`FitManager`	Unified training manager for Estimator and Parametric models in ensembles.

Classes#

`FitManager(default_trainer: pl.Trainer | None = None, default_trainer_kwargs: dict | None = None, default_datamodule_kwargs: dict | None = None, estimator_trainers: dict[str, pl.Trainer] | None = None, estimator_trainer_kwargs: dict[str, dict] | None = None, estimator_datamodule_kwargs: dict[str, dict] | None = None)` #

Unified training manager for Estimator and Parametric models in ensembles.

This class provides a separation between ensemble algorithms model training. It handles the different training requirements of non-parametric methods (direct fit) and parametric neural network methods (PyTorch Lightning Trainer).

PARAMETER	DESCRIPTION
`default_trainer`	Default PyTorch Lightning Trainer instance to use for all Parametric models. If None, a Trainer will be created from `default_trainer_kwargs` when needed. TYPE: `pl.Trainer \| None, optional, by default None` DEFAULT: `None`
`default_trainer_kwargs`	Default keyword arguments for creating PyTorch Lightning Trainers when `default_trainer` is not provided. `max_epochs`: Maximum number of training epochs `callbacks`: List of Lightning callbacks (EarlyStopping, etc.) `logger`: Lightning logger `accelerator`: Hardware accelerator ('gpu', 'cpu') `devices`: Number of devices to use `precision`: Training precision TYPE: `dict \| None, optional, by default None` DEFAULT: `None`
`default_datamodule_kwargs`	Default keyword arguments for creating DataModuleWeighted instances. `batch_size`: Batch size for training `val`: Validation split fraction (0.0-1.0) `use_random_split`: Whether to use random or sequential validation split `shuffle`: Whether to shuffle training data TYPE: `dict \| None, optional, by default None` DEFAULT: `None`
`estimator_trainers`	Dictionary mapping estimator names to specific Trainer instances. This allows different estimators in the same ensemble to use different Trainers with distinct configurations (e.g., different callbacks, loggers). TYPE: `dict[str, pl.Trainer] \| None, optional, by default None` DEFAULT: `None`
`estimator_trainer_kwargs`	Dictionary mapping estimator names to Trainer kwargs. Used to create per-estimator Trainers that override the default configuration. TYPE: `dict[str, dict] \| None, optional, by default None` DEFAULT: `None`
`estimator_datamodule_kwargs`	Dictionary mapping estimator names to DataModule kwargs. Allows different batch sizes, validation splits, etc., for different estimators. TYPE: `dict[str, dict] \| None, optional, by default None` DEFAULT: `None`

Examples:

Basic usage with default Trainer for all Parametric models:

>>> import pytorch_lightning as pl
>>> from spectre.ensemble import FitManager, BootstrapSpectral
>>> from spectre.parametric import SpectralMap
>>>
>>> # Configure default Trainer
>>> trainer = pl.Trainer(max_epochs=50, accelerator="gpu")
>>> fit_manager = FitManager(default_trainer=trainer)
>>>
>>> # Create ensemble with Parametric model
>>> sm = SpectralMap(model=network)
>>> bootstrap = BootstrapSpectral(
...     estimator=sm,
...     n_estimators=10,
...     fit_manager=fit_manager,
... )
>>> bootstrap.fit(X)

Per-estimator Trainer overrides for mixed ensembles:

>>> from spectre.decomposition import DiffusionMap
>>> from spectre.ensemble import AdaptiveWeightedEnsemble
>>>
>>> # Default trainer for most models
>>> default_trainer = pl.Trainer(max_epochs=50)
>>>
>>> # Special trainer for one specific model
>>> special_trainer = pl.Trainer(
...     max_epochs=100,
...     callbacks=[pl.callbacks.EarlyStopping(monitor="val_loss")],
... )
>>>
>>> fit_manager = FitManager(
...     default_trainer=default_trainer,
...     estimator_trainers={"spectral_map": special_trainer},
... )
>>>
>>> # Mixed ensemble: non-parametric + parametric
>>> dm = DiffusionMap(n_components=10)
>>> sm = SpectralMap(model=network)
>>>
>>> ensemble = AdaptiveWeightedEnsemble(
...     estimators={"dm": dm, "spectral_map": sm}, fit_manager=fit_manager
... )
>>> ensemble.fit(X)  # dm uses direct fit(), sm uses special_trainer

Using trainer kwargs instead of instances:

>>> fit_manager = FitManager(
...     default_trainer_kwargs={"max_epochs": 50, "accelerator": "gpu"},
...     estimator_trainer_kwargs={
...         "model_a": {"max_epochs": 100},
...         "model_b": {"max_epochs": 25},
...     },
... )

Notes

Estimator (non-parametric) models are trained via direct fit() call
Parametric (neural network) models use PyTorch Lightning Trainer

METHOD	DESCRIPTION
`fit`	Fit an estimator using the appropriate training strategy.

Source code in spectre/ensemble/fit_manager.py

def __init__(
    self,
    default_trainer: pl.Trainer | None = None,
    default_trainer_kwargs: dict | None = None,
    default_datamodule_kwargs: dict | None = None,
    estimator_trainers: dict[str, pl.Trainer] | None = None,
    estimator_trainer_kwargs: dict[str, dict] | None = None,
    estimator_datamodule_kwargs: dict[str, dict] | None = None,
) -> None:
    self.default_trainer = default_trainer
    self.default_trainer_kwargs = default_trainer_kwargs or {}
    self.default_datamodule_kwargs = default_datamodule_kwargs or {}
    self.estimator_trainers = estimator_trainers or {}
    self.estimator_trainer_kwargs = estimator_trainer_kwargs or {}
    self.estimator_datamodule_kwargs = estimator_datamodule_kwargs or {}

Functions#

`fit(estimator: ModelType, X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None, estimator_name: str | None = None) -> ModelType` #

Fit an estimator using the appropriate training strategy.

Detects whether the estimator is a non-parametric Estimator or a Parametric model and applies the correct training approach.

PARAMETER	DESCRIPTION
`estimator`	Model to train. Can be either a non-parametric Estimator or a Parametric model. TYPE: `Estimator \| Parametric`
`X`	Training data of shape (n_samples, n_features). TYPE: `Tensor`
`weights`	Sample weights of shape (n_samples,). TYPE: `torch.Tensor \| None, optional, by default None` DEFAULT: `None`
`target`	Target values for supervised learning. TYPE: `torch.Tensor \| None, optional, by default None` DEFAULT: `None`
`estimator_name`	Name of the estimator for looking up per-estimator configuration overrides. If provided and found in `estimator_trainers` or `estimator_trainer_kwargs`, uses those settings instead of defaults. TYPE: `str \| None, optional, by default None` DEFAULT: `None`

RETURNS	DESCRIPTION
`Estimator \| Parametric`	The fitted estimator (same instance that was passed in).