core.posterior_models.base_model.BasePosteriorModel
core.posterior_models.base_model.BasePosteriorModel(
model_filename=None,
metadata=None,
initial_weights=None,
device='cuda',
load_training_info=True,
)Abstract base class for PosteriorModels. This is intended to construct and hold a neural network for estimating the posterior density, as well as saving / loading, and training.
Subclasses must implement methods for constructing the specific network, sampling, density evaluation, and computing the loss during training.
Initialize a model for the posterior distribution.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| model_filename | str | If given, loads data from the given file. | None |
| metadata | dict | If given, initializes the model from these settings | None |
| initial_weights | dict | Initial weights for the model | None |
| device | str | 'cuda' |
|
| load_training_info | bool | True |
Attributes
| Name | Description |
|---|---|
| context | |
| device | |
| epoch | |
| event_metadata | |
| initial_weights | |
| metadata | |
| model_kwargs | |
| network | |
| network_kwargs | |
| optimizer | |
| optimizer_kwargs | |
| scheduler | |
| scheduler_kwargs | |
| version |
Methods
| Name | Description |
|---|---|
| initialize_network | Initialize the network backbone for the posterior model. |
| initialize_optimizer_and_scheduler | Initializes the optimizer and scheduler with self.optimizer_kwargs |
| load_model | Load a posterior model from the disk. |
| log_prob | Evaluate the log posterior density, |
| loss | Compute the loss for a batch of data. |
| network_to_device | Put model to device, and set self.device accordingly. |
| sample | Sample parameters theta from the posterior model, |
| sample_and_log_prob | Sample parameters theta from the posterior model, |
| save_model | Save the posterior model to the disk. |
| train |
initialize_network
core.posterior_models.base_model.BasePosteriorModel.initialize_network()Initialize the network backbone for the posterior model.
initialize_optimizer_and_scheduler
core.posterior_models.base_model.BasePosteriorModel.initialize_optimizer_and_scheduler(
)Initializes the optimizer and scheduler with self.optimizer_kwargs and self.scheduler_kwargs, respectively.
load_model
core.posterior_models.base_model.BasePosteriorModel.load_model(
model_filename,
load_training_info=True,
device='cuda',
)Load a posterior model from the disk.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| model_filename | str | path to saved model | required |
| load_training_info | bool | specifies whether information required to proceed with training is loaded, e.g. optimizer state dict | True |
| device | str | 'cuda' |
log_prob
core.posterior_models.base_model.BasePosteriorModel.log_prob(theta, *context)Evaluate the log posterior density,
log p(theta | context)
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| theta | torch.Tensor | Parameter values at which to evaluate the density. Should have a batch dimension (even if size B = 1). | required |
| context | torch.Tensor | Context information (typically observed data). Must have context.shape[0] = B. | () |
Returns
| Name | Type | Description |
|---|---|---|
| log_prob | torch.Tensor | Shape (B,) |
loss
core.posterior_models.base_model.BasePosteriorModel.loss(theta, *context)Compute the loss for a batch of data.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| theta | torch.Tensor | Parameter values at which to evaluate the density. Should have a batch dimension (even if size B = 1). | required |
| context | torch.Tensor | Context information (typically observed data). Must have the same leading (batch) dimension as theta. | () |
Returns
| Name | Type | Description |
|---|---|---|
| loss | torch.Tensor | Mean loss across the batch (a scalar). |
network_to_device
core.posterior_models.base_model.BasePosteriorModel.network_to_device(device)Put model to device, and set self.device accordingly.
sample
core.posterior_models.base_model.BasePosteriorModel.sample(
*context,
num_samples=1,
)Sample parameters theta from the posterior model,
theta ~ p(theta | context)
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| context | torch.Tensor | Context information (typically observed data). Should have a batch dimension (even if size B = 1). | () |
| num_samples | int | Number of samples to generate. | 1 |
Returns
| Name | Type | Description |
|---|---|---|
| samples | torch.Tensor | Shape (B, num_samples, dim(theta)) |
sample_and_log_prob
core.posterior_models.base_model.BasePosteriorModel.sample_and_log_prob(
*context,
num_samples=1,
)Sample parameters theta from the posterior model,
theta ~ p(theta | context)
and also return the log_prob. For models such as normalizing flows, it is more economical to calculate the log_prob at the same time as sampling, rather than as a separate step.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| context | torch.Tensor | Context information (typically observed data). Should have a batch dimension (even if size B = 1). | () |
| num_samples | int | Number of samples to generate. | 1 |
Returns
| Name | Type | Description |
|---|---|---|
| samples, log_prob: torch.Tensor, torch.Tensor | Shapes (B, num_samples, dim(theta)), (B, num_samples) |
save_model
core.posterior_models.base_model.BasePosteriorModel.save_model(
model_filename,
save_training_info=True,
)Save the posterior model to the disk.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| model_filename | str | filename for saving the model | required |
| save_training_info | bool | specifies whether information required to proceed with training is saved, e.g. optimizer state dict | True |
train
core.posterior_models.base_model.BasePosteriorModel.train(
train_loader,
test_loader,
train_dir,
runtime_limits=None,
checkpoint_epochs=None,
use_wandb=False,
test_only=False,
early_stopping=None,
)Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| train_loader | torch.utils.data.DataLoader | required | |
| test_loader | torch.utils.data.DataLoader | required | |
| train_dir | str | required | |
| runtime_limits | object | None |
|
| checkpoint_epochs | int | None |
|
| use_wandb | False |
||
| test_only | if True, training is skipped | False |
|
| early_stopping | Optional[EarlyStopping] | Optional EarlyStopping instance. | None |