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
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
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
- 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
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