pyrelational.model_managers¶

Abstract Model Manager¶

class ModelManager(model_class: Type[ModelType], model_config: str | Dict[str, Any], trainer_config: str | Dict[str, Any])[source]¶

Bases: ABC, Generic[ModelType, E]

Abstract class used to wrap models to interact with the Strategy. It handles model instantiation at each iteration, training, testing, and queries.

Parameters:

model_class – a model constructor (e.g. torch.nn.Linear)
model_config – a dictionary containing the config required to instantiate a model form the model_class (e.g. {in_features=100, out_features=34, bias=True, device=None, dtype=None} for a torch.nn.Linear constructor)
trainer_config – a dictionary containing the config required to instantiate the trainer module/function

__call__(loader: DataLoader[Any]) → Any[source]¶

Call method to output model predictions

Parameters:: loader – pytorch dataloader
Returns:: model predictions for each sample in dataloader

is_trained() → bool[source]¶: Check if model was trained.

reset() → None[source]¶: Reset stored _current_model.

abstract test(loader: DataLoader[Any]) → Dict[str, float][source]¶

Run test routine.

Parameters:: loader – pytorch dataloader for test set
Returns:: performance metrics

abstract train(train_loader: DataLoader[Any], valid_loader: DataLoader[Any] | None = None) → None[source]¶

Run train routine.

Parameters:

train_loader – pytorch dataloader for training set
valid_loader – pytorch dataloader for validation set

Pytorch Lightning Model¶

class LightningModelManager(model_class: Type[LightningModule], model_config: Dict[str, Any] | str, trainer_config: Dict[str, Any] | str)[source]¶

Bases: ModelManager[LightningModule, LightningModule]

A wrapper for pytorch lightning modules that instantiates and uses a pytorch lightning trainer.

Example:

import torch
import lightning.pytorch as pl

class PyLModel(pl.LightningModule):
      def __init__(self, in_dim, out_dim):
          super(PyLModel, self).()
          self.linear = torch.nn.Linear(in_dim, out_dim)
   # need to define other train/test steps and optimizers methods required
   # by pytorch-lightning to run this example

wrapper = LightningModelManager(
                PyLModel,
                model_config={"in_dim":10, "out_dim":1},
                trainer_config={"epochs":100},
          )
wrapper.train(train_loader, valid_loader)

Parameters:

model_class – a model constructor class which inherits from pytorch lightning (see above example)
model_config – a dictionary containing the config required to instantiate a model form the model_class (e.g. see above example)
trainer_config – a dictionary containing the config required to instantiate the pytorch lightning trainer

__call__(loader: DataLoader[Any]) → Tensor[source]¶

Call function which outputs model predictions from dataloader

Parameters:: loader – pytorch dataloader
Returns:: model predictions of shape (number of samples in loader,1)

init_trainer() → Tuple[Trainer, ModelCheckpoint][source]¶

Initialise pytorch lightning trainer.

Returns:: a pytorch lightning trainer object

test(loader: DataLoader[Any]) → Dict[str, float][source]¶

Run test routine.

Parameters:: loader – pytorch dataloader for test set
Returns:: performance metrics

train(train_loader: DataLoader[Any], valid_loader: DataLoader[Any] | None = None) → None[source]¶

Run train routine.

Parameters:

train_loader – pytorch dataloader for training set
valid_loader – pytorch dataloader for validation set

Ensemble Models¶

class EnsembleModelManager(model_class: Type[ModelType], model_config: str | Dict[str, Any], trainer_config: str | Dict[str, Any], n_estimators: int = 10)[source]¶

Bases: Generic[ModelType], ModelManager[ModelType, List[ModelType]], ABC

Generic wrapper for ensemble uncertainty estimator

Parameters:

model_class – a model constructor (e.g. torch.nn.Linear)
model_config – a dictionary containing the config required to instantiate a model form the model_class (e.g. {in_features=100, out_features=34, bias=True, device=None, dtype=None} for a torch.nn.Linear constructor)
trainer_config – a dictionary containing the config required to instantiate the trainer module/function
n_estimators – number of models in ensemble

__call__(loader: DataLoader[Any]) → Tensor[source]¶

Call method to output model predictions for each model in the ensemble

Parameters:: loader – pytorch dataloader
Returns:: model predictions of shape (n_estimators, number of samples in loader, 1)

class LightningEnsembleModelManager(model_class: Type[LightningModule], model_config: Dict[str, Any] | str, trainer_config: Dict[str, Any] | str, n_estimators: int = 10)[source]¶

Bases: EnsembleModelManager[LightningModule], LightningModelManager

Wrapper for ensemble estimator with pytorch lightning trainer

Example:

import torch
import lightning.pytorch as pl

class PyLModel(pl.LightningModule):
   def __init__(self, in_dim, out_dim):
       super(PyLModel, self).()
       self.linear = torch.nn.Linear(in_dim, out_dim)
# need to define other train/test steps and optimizers methods required
# by pytorch-lightning to run this example

wrapper = LightningEnsembleModelManager(
             PyLModel,
             model_config={"in_dim":10, "out_dim":1},
             trainer_config={"epochs":100},
             n_estimators=10,
       )
wrapper.train(train_loader, valid_loader)
predictions = wrapper(loader)
assert predictions.size(0) == 10

Parameters:

model_class – a model constructor class which inherits from pytorch lightning (see above example)
model_config – a dictionary containing the config required to instantiate a model form the model_class (e.g. see above example)
trainer_config – a dictionary containing the config required to instantiate the pytorch lightning trainer
n_estimators – number of models in ensemble

test(loader: DataLoader[Any]) → Dict[str, float][source]¶

Test ensemble model. The mean performance across all the models in the ensemble is reported for each metric

Parameters:: loader – dataloader for test set
Returns:: average performance for each metric (defined in the model_class)

train(train_loader: DataLoader[Any], valid_loader: DataLoader[Any] | None = None) → None[source]¶

Train each model in ensemble.

Parameters:

train_loader – pytorch data loader containing train data
valid_loader – pytorch data loader containing validation data

MCDropout Models¶

class LightningMCDropoutModelManager(model_class: Type[LightningModule], model_config: Dict[str, Any] | str, trainer_config: Dict[str, Any] | str, n_estimators: int = 10, eval_dropout_prob: float = 0.2)[source]¶

Bases: MCDropoutModelManager, LightningModelManager

Wrapper for MC Dropout estimator with pytorch lightning trainer

Example:

import torch
import lightning.pytorch as pl

class PyLModel(pl.LightningModule):
      def __init__(self, in_dim, out_dim):
          super(PyLModel, self).()
          self.linear = torch.nn.Linear(in_dim, out_dim)
   # need to define other train/test steps and optimizers methods required
   # by pytorch-lightning to run this example

wrapper = LightningMCDropoutModelManager(
                PyLModel,
                model_config={"in_dim":10, "out_dim":1},
                trainer_config={"epochs":100},
                n_estimators=10,
                eval_dropout_prob=0.2,
          )
wrapper.train(train_loader, valid_loader)
predictions = wrapper(loader)
assert predictions.size(0) == 10

Parameters:

model_class – a model constructor class which inherits from pytorch lightning (see above example)
model_config – a dictionary containing the config required to instantiate a model form the model_class (e.g. see above example)
trainer_config – a dictionary containing the config required to instantiate the pytorch lightning trainer
n_estimators – number of times to sample a prediction for each input
eval_dropout_prob – dropout parameter used when accessing model predictions

class MCDropoutModelManager(model_class: Type[Module], model_config: str | Dict[str, Any], trainer_config: str | Dict[str, Any], n_estimators: int = 10, eval_dropout_prob: float = 0.2)[source]¶

Bases: ModelManager[Module, Module], ABC

Generic model wrapper for mcdropout uncertainty estimator

Parameters:

model_class – a model constructor (e.g. torch.nn.Linear)
model_config – a dictionary containing the config required to instantiate a model form the model_class (e.g. {in_features=100, out_features=34, bias=True, device=None, dtype=None} for a torch.nn.Linear constructor)
trainer_config – a dictionary containing the config required to instantiate the trainer module/function
n_estimators – number of times to sample a prediction for each input
eval_dropout_prob – dropout parameter used when accessing model predictions

__call__(loader: DataLoader[Any]) → Tensor[source]¶

Call function which outputs model predictions using dropout

Parameters:: loader – pytorch dataloader
Returns:: model predictions of shape (n_estimators, number of samples in loader, 1)