def __init__(
self,
method_family: str,
neural_net: nn.Module,
prior,
x_shape: torch.Size,
mcmc_method: str = "slice_np",
mcmc_parameters: Optional[Dict[str, Any]] = None,
device: str = "cpu",
):
"""
Args:
method_family: One of snpe, snl, snre_a or snre_b.
neural_net: A classifier for SNRE, a density estimator for SNPE and SNL.
prior: Prior distribution with `.log_prob()` and `.sample()`.
x_shape: Shape of a single simulator output.
mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`,
`hmc`, `nuts`. Currently defaults to `slice_np` for a custom numpy
implementation of slice sampling; select `hmc`, `nuts` or `slice` for
Pyro-based sampling.
mcmc_parameters: Dictionary overriding the default parameters for MCMC.
The following parameters are supported: `thin` to set the thinning
factor for the chain, `warmup_steps` to set the initial number of
samples to discard, `num_chains` for the number of chains,
`init_strategy` for the initialisation strategy for chains; `prior`
will draw init locations from prior, whereas `sir` will use Sequential-
Importance-Resampling using `init_strategy_num_candidates` to find init
locations.
device: Training device, e.g., cpu or cuda:0.
"""
kwargs = del_entries(locals(), entries=("self", "__class__"))
super().__init__(**kwargs)
def __call__(
self,
num_rounds: int,
num_simulations_per_round: OneOrMore[int],
x_o: Optional[Tensor] = None,
training_batch_size: int = 50,
learning_rate: float = 5e-4,
validation_fraction: float = 0.1,
stop_after_epochs: int = 20,
max_num_epochs: Optional[int] = None,
clip_max_norm: Optional[float] = 5.0,
exclude_invalid_x: bool = True,
discard_prior_samples: bool = False,
retrain_from_scratch_each_round: bool = False,
) -> NeuralPosterior:
r"""Run AALR / SNRE_A.
Return posterior $p(\theta|x)$ after inference (possibly over several rounds).
Args:
num_rounds: Number of rounds to run. Each round consists of a simulation and
training phase. `num_rounds=1` leads to a posterior $p(\theta|x)$ valid
for _any_ $x$ (amortized), but requires many simulations.
Alternatively, with `num_rounds>1` the inference returns a posterior
$p(\theta|x_o)$ focused on a specific observation `x_o`, potentially
requiring less simulations.
num_simulations_per_round: Number of simulator calls per round.
x_o: An observation that is only required when doing inference
over multiple rounds. After the first round, `x_o` is used to guide the
sampling so that the simulator is run with parameters that are likely
for that `x_o`, i.e. they are sampled from the posterior obtained in the
previous round $p(\theta|x_o)$.
training_batch_size: Training batch size.
learning_rate: Learning rate for Adam optimizer.
validation_fraction: The fraction of data to use for validation.
stop_after_epochs: The number of epochs to wait for improvement on the
validation set before terminating training.
max_num_epochs: Maximum number of epochs to run. If reached, we stop
training even when the validation loss is still decreasing. If None, we
train until validation loss increases (see also `stop_after_epochs`).
clip_max_norm: Value at which to clip the total gradient norm in order to
prevent exploding gradients. Use None for no clipping.
exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
during training. Expect errors, silent or explicit, when `False`.
discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
from the prior. Training may be sped up by ignoring such less targeted
samples.
retrain_from_scratch_each_round: Whether to retrain the conditional density
estimator for the posterior from scratch each round.
Returns:
Posterior $p(\theta|x)$ that can be sampled and evaluated.
"""
# AALR is defined for `num_atoms=2`.
# Proxy to `super().__call__` to ensure right parameter.
kwargs = del_entries(locals(), entries=("self", "__class__"))
return super().__call__(**kwargs, num_atoms=2)
def __init__(
self,
method_family: str,
neural_net: nn.Module,
prior,
x_shape: torch.Size,
rejection_sampling_parameters: Optional[Dict[str, Any]] = None,
sample_with_mcmc: bool = True,
mcmc_method: str = "slice_np",
mcmc_parameters: Optional[Dict[str, Any]] = None,
device: str = "cpu",
):
"""
Args:
method_family: One of snpe, snl, snre_a or snre_b.
neural_net: A classifier for SNRE, a density estimator for SNPE and SNL.
prior: Prior distribution with `.log_prob()` and `.sample()`.
x_shape: Shape of a single simulator output.
rejection_sampling_parameters: Dictonary overriding the default parameters
for rejection sampling. The following parameters are supported:
`max_sampling_batch_size` to set the batch size for drawing new
samples from the candidate distribution, e.g., the posterior. Larger
batch size speeds up sampling.
sample_with_mcmc: Whether to sample with MCMC. Will always be `True` for SRE
and SNL, but can also be set to `True` for SNPE if MCMC is preferred to
deal with leakage over rejection sampling.
mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`,
`hmc`, `nuts`. Currently defaults to `slice_np` for a custom numpy
implementation of slice sampling; select `hmc`, `nuts` or `slice` for
Pyro-based sampling.
mcmc_parameters: Dictionary overriding the default parameters for MCMC.
The following parameters are supported: `thin` to set the thinning
factor for the chain, `warmup_steps` to set the initial number of
samples to discard, `num_chains` for the number of chains,
`init_strategy` for the initialisation strategy for chains; `prior`
will draw init locations from prior, whereas `sir` will use Sequential-
Importance-Resampling using `init_strategy_num_candidates` to find init
locations.
device: Training device, e.g., cpu or cuda:0
"""
kwargs = del_entries(
locals(),
entries=(
"self",
"__class__",
"sample_with_mcmc",
"rejection_sampling_parameters",
),
)
super().__init__(**kwargs)
self.set_sample_with_mcmc(sample_with_mcmc)
self.set_rejection_sampling_parameters(rejection_sampling_parameters)
self._purpose = (
"It allows to .sample() and .log_prob() the posterior and wraps the "
"output of the .net to avoid leakage into regions with 0 prior probability."
)
def __init__(
self,
prior: Optional[Distribution] = None,
density_estimator: Union[str, Callable] = "maf",
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[TensorboardSummaryWriter] = None,
show_progress_bars: bool = True,
):
r"""SNPE-C / APT [1].
[1] _Automatic Posterior Transformation for Likelihood-free Inference_,
Greenberg et al., ICML 2019, https://arxiv.org/abs/1905.07488.
This class implements two loss variants of SNPE-C: the non-atomic and the atomic
version. The atomic loss of SNPE-C can be used for any density estimator,
i.e. also for normalizing flows. However, it suffers from leakage issues. On
the other hand, the non-atomic loss can only be used only if the proposal
distribution is a mixture of Gaussians, the density estimator is a mixture of
Gaussians, and the prior is either Gaussian or Uniform. It does not suffer from
leakage issues. At the beginning of each round, we print whether the non-atomic
or the atomic version is used.
In this codebase, we will automatically switch to the non-atomic loss if the
following criteria are fulfilled:<br/>
- proposal is a `DirectPosterior` with density_estimator `mdn`, as built
with `utils.sbi.posterior_nn()`.<br/>
- the density estimator is a `mdn`, as built with
`utils.sbi.posterior_nn()`.<br/>
- `isinstance(prior, MultivariateNormal)` (from `torch.distributions`) or
`isinstance(prior, sbi.utils.BoxUniform)`
Note that custom implementations of any of these densities (or estimators) will
not trigger the non-atomic loss, and the algorithm will fall back onto using
the atomic loss.
Args:
prior: A probability distribution that expresses prior knowledge about the
parameters, e.g. which ranges are meaningful for them.
density_estimator: If it is a string, use a pre-configured network of the
provided type (one of nsf, maf, mdn, made). Alternatively, a function
that builds a custom neural network can be provided. The function will
be called with the first batch of simulations (theta, x), which can
thus be used for shape inference and potentially for z-scoring. It
needs to return a PyTorch `nn.Module` implementing the density
estimator. The density estimator needs to provide the methods
`.log_prob` and `.sample()`.
device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
logging_level: Minimum severity of messages to log. One of the strings
INFO, WARNING, DEBUG, ERROR and CRITICAL.
summary_writer: A tensorboard `SummaryWriter` to control, among others, log
file location (default is `<current working directory>/logs`.)
show_progress_bars: Whether to show a progressbar during training.
"""
kwargs = del_entries(locals(), entries=("self", "__class__"))
super().__init__(**kwargs)
def __init__(
self,
method_family: str,
neural_net: nn.Module,
prior,
x_shape: torch.Size,
sample_with: str = "rejection",
mcmc_method: str = "slice_np",
mcmc_parameters: Optional[Dict[str, Any]] = None,
rejection_sampling_parameters: Optional[Dict[str, Any]] = None,
device: str = "cpu",
):
"""
Args:
method_family: One of snpe, snl, snre_a or snre_b.
neural_net: A classifier for SNRE, a density estimator for SNPE and SNL.
prior: Prior distribution with `.log_prob()` and `.sample()`.
x_shape: Shape of a single simulator output.
sample_with: Method to use for sampling from the posterior. Must be one of
[`mcmc` | `rejection`]. With default parameters, `rejection` samples
from the posterior estimated by the neural net and rejects only if the
samples are outside of the prior support.
mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`,
`hmc`, `nuts`. Currently defaults to `slice_np` for a custom numpy
implementation of slice sampling; select `hmc`, `nuts` or `slice` for
Pyro-based sampling.
mcmc_parameters: Dictionary overriding the default parameters for MCMC.
The following parameters are supported: `thin` to set the thinning
factor for the chain, `warmup_steps` to set the initial number of
samples to discard, `num_chains` for the number of chains,
`init_strategy` for the initialisation strategy for chains; `prior`
will draw init locations from prior, whereas `sir` will use Sequential-
Importance-Resampling using `init_strategy_num_candidates` to find init
locations.
rejection_sampling_parameters: Dictionary overriding the default parameters
for rejection sampling. The following parameters are supported:
`proposal` as the proposal distribtution (default is the trained
neural net). `max_sampling_batch_size` as the batchsize of samples
being drawn from the proposal at every iteration.
`num_samples_to_find_max` as the number of samples that are used to
find the maximum of the `potential_fn / proposal` ratio.
`num_iter_to_find_max` as the number of gradient ascent iterations to
find the maximum of that ratio. `m` as multiplier to that ratio.
device: Training device, e.g., cpu or cuda:0
"""
kwargs = del_entries(
locals(),
entries=("self", "__class__"),
)
super().__init__(**kwargs)
self._purpose = (
"It allows to .sample() and .log_prob() the posterior and wraps the "
"output of the .net to avoid leakage into regions with 0 prior probability."
)
def __init__(
self,
method_family: str,
neural_net: nn.Module,
prior,
x_shape: torch.Size,
sample_with: str = "mcmc",
mcmc_method: str = "slice_np",
mcmc_parameters: Optional[Dict[str, Any]] = None,
rejection_sampling_parameters: Optional[Dict[str, Any]] = None,
device: str = "cpu",
):
"""
Args:
method_family: One of snpe, snl, snre_a or snre_b.
neural_net: A classifier for SNRE, a density estimator for SNPE and SNL.
prior: Prior distribution with `.log_prob()` and `.sample()`.
x_shape: Shape of the simulated data. It can differ from the
observed data the posterior is conditioned on later in the batch
dimension. If it differs, the additional entries are interpreted as
independent and identically distributed data / trials. I.e., the data is
assumed to be generated based on the same (unknown) model parameters or
experimental condations.
sample_with: Method to use for sampling from the posterior. Must be one of
[`mcmc` | `rejection`].
mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`,
`hmc`, `nuts`. Currently defaults to `slice_np` for a custom numpy
implementation of slice sampling; select `hmc`, `nuts` or `slice` for
Pyro-based sampling.
mcmc_parameters: Dictionary overriding the default parameters for MCMC.
The following parameters are supported: `thin` to set the thinning
factor for the chain, `warmup_steps` to set the initial number of
samples to discard, `num_chains` for the number of chains,
`init_strategy` for the initialisation strategy for chains; `prior`
will draw init locations from prior, whereas `sir` will use Sequential-
Importance-Resampling using `init_strategy_num_candidates` to find init
locations.
rejection_sampling_parameters: Dictionary overriding the default parameters
for rejection sampling. The following parameters are supported:
`proposal` as the proposal distribtution (default is the prior).
`max_sampling_batch_size` as the batchsize of samples being drawn from
the proposal at every iteration. `num_samples_to_find_max` as the
number of samples that are used to find the maximum of the
`potential_fn / proposal` ratio. `num_iter_to_find_max` as the number
of gradient ascent iterations to find the maximum of that ratio. `m` as
multiplier to that ratio.
device: Training device, e.g., cpu or cuda:0.
"""
kwargs = del_entries(locals(), entries=("self", "__class__"))
super().__init__(**kwargs)
self._purpose = (
"It provides MCMC to .sample() from the posterior and "
"can evaluate the _unnormalized_ posterior density with .log_prob()."
)
def train(
self,
training_batch_size: int = 50,
learning_rate: float = 5e-4,
validation_fraction: float = 0.1,
stop_after_epochs: int = 20,
max_num_epochs: Optional[int] = None,
clip_max_norm: Optional[float] = 5.0,
exclude_invalid_x: bool = True,
resume_training: bool = False,
discard_prior_samples: bool = False,
retrain_from_scratch_each_round: bool = False,
show_train_summary: bool = False,
dataloader_kwargs: Optional[Dict] = None,
) -> NeuralPosterior:
r"""
Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
Args:
training_batch_size: Training batch size.
learning_rate: Learning rate for Adam optimizer.
validation_fraction: The fraction of data to use for validation.
stop_after_epochs: The number of epochs to wait for improvement on the
validation set before terminating training.
max_num_epochs: Maximum number of epochs to run. If reached, we stop
training even when the validation loss is still decreasing. If None, we
train until validation loss increases (see also `stop_after_epochs`).
clip_max_norm: Value at which to clip the total gradient norm in order to
prevent exploding gradients. Use None for no clipping.
exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
during training. Expect errors, silent or explicit, when `False`.
resume_training: Can be used in case training time is limited, e.g. on a
cluster. If `True`, the split between train and validation set, the
optimizer, the number of epochs, and the best validation log-prob will
be restored from the last time `.train()` was called.
discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
from the prior. Training may be sped up by ignoring such less targeted
samples.
retrain_from_scratch_each_round: Whether to retrain the conditional density
estimator for the posterior from scratch each round.
show_train_summary: Whether to print the number of epochs and validation
loss and leakage after the training.
dataloader_kwargs: Additional or updated kwargs to be passed to the training
and validation dataloaders (like, e.g., a collate_fn)
Returns:
Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
"""
# AALR is defined for `num_atoms=2`.
# Proxy to `super().__call__` to ensure right parameter.
kwargs = del_entries(locals(), entries=("self", "__class__"))
return super().train(**kwargs, num_atoms=2)
def __call__(
self,
num_simulations: int,
proposal: Optional[Any] = None,
training_batch_size: int = 50,
learning_rate: float = 5e-4,
validation_fraction: float = 0.1,
stop_after_epochs: int = 20,
max_num_epochs: Optional[int] = None,
clip_max_norm: Optional[float] = 5.0,
exclude_invalid_x: bool = True,
discard_prior_samples: bool = False,
retrain_from_scratch_each_round: bool = False,
device='cpu'
) -> NeuralPosterior:
r"""Run AALR / SNRE_A.
Return posterior $p(\theta|x)$ after inference.
Args:
num_simulations: Number of simulator calls.
proposal: Distribution that the parameters $\theta$ are drawn from.
`proposal=None` uses the prior. Setting the proposal to a distribution
targeted on a specific observation, e.g. a posterior $p(\theta|x_o)$
obtained previously, can lead to less required simulations.
training_batch_size: Training batch size.
learning_rate: Learning rate for Adam optimizer.
validation_fraction: The fraction of data to use for validation.
stop_after_epochs: The number of epochs to wait for improvement on the
validation set before terminating training.
max_num_epochs: Maximum number of epochs to run. If reached, we stop
training even when the validation loss is still decreasing. If None, we
train until validation loss increases (see also `stop_after_epochs`).
clip_max_norm: Value at which to clip the total gradient norm in order to
prevent exploding gradients. Use None for no clipping.
exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
during training. Expect errors, silent or explicit, when `False`.
discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
from the prior. Training may be sped up by ignoring such less targeted
samples.
retrain_from_scratch_each_round: Whether to retrain the conditional density
estimator for the posterior from scratch each round.
Returns:
Posterior $p(\theta|x)$ that can be sampled and evaluated.
"""
# AALR is defined for `num_atoms=2`.
# Proxy to `super().__call__` to ensure right parameter.
kwargs = del_entries(locals(), entries=("self", "__class__"))
return super().__call__(**kwargs, num_atoms=2)
def __init__(
self,
prior: Optional[Distribution] = None,
density_estimator: Union[str, Callable] = "mnle",
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[TensorboardSummaryWriter] = None,
show_progress_bars: bool = True,
):
r"""Mixed Neural Likelihood Estimation (MNLE) [1].
Like SNLE, but not sequential and designed to be applied to data with mixed
types, e.g., continuous data and discrete data like they occur in
decision-making experiments (reation times and choices).
[1] Flexible and efficient simulation-based inference for models of
decision-making, Boelts et al. 2021,
https://www.biorxiv.org/content/10.1101/2021.12.22.473472v2
Args:
prior: A probability distribution that expresses prior knowledge about the
parameters, e.g. which ranges are meaningful for them. If `None`, the
prior must be passed to `.build_posterior()`.
density_estimator: If it is a string, it must be "mnle" to use the
preconfiugred neural nets for MNLE. Alternatively, a function
that builds a custom neural network can be provided. The function will
be called with the first batch of simulations (theta, x), which can
thus be used for shape inference and potentially for z-scoring. It
needs to return a PyTorch `nn.Module` implementing the density
estimator. The density estimator needs to provide the methods
`.log_prob`, `.log_prob_iid()` and `.sample()`.
device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
logging_level: Minimum severity of messages to log. One of the strings
INFO, WARNING, DEBUG, ERROR and CRITICAL.
summary_writer: A tensorboard `SummaryWriter` to control, among others, log
file location (default is `<current working directory>/logs`.)
show_progress_bars: Whether to show a progressbar during simulation and
sampling.
"""
if isinstance(density_estimator, str):
assert (
density_estimator == "mnle"
), f"""MNLE can be used with preconfigured 'mnle' density estimator only,
not with {density_estimator}."""
kwargs = del_entries(locals(), entries=("self", "__class__"))
super().__init__(**kwargs)
def train(
self,
num_atoms: int = 10,
training_batch_size: int = 50,
learning_rate: float = 5e-4,
validation_fraction: float = 0.1,
stop_after_epochs: int = 20,
max_num_epochs: Optional[int] = None,
clip_max_norm: Optional[float] = 5.0,
exclude_invalid_x: bool = True,
discard_prior_samples: bool = False,
retrain_from_scratch_each_round: bool = False,
show_train_summary: bool = False,
) -> NeuralPosterior:
r"""
Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
Args:
num_atoms: Number of atoms to use for classification.
training_batch_size: Training batch size.
learning_rate: Learning rate for Adam optimizer.
validation_fraction: The fraction of data to use for validation.
stop_after_epochs: The number of epochs to wait for improvement on the
validation set before terminating training.
max_num_epochs: Maximum number of epochs to run. If reached, we stop
training even when the validation loss is still decreasing. If None, we
train until validation loss increases (see also `stop_after_epochs`).
clip_max_norm: Value at which to clip the total gradient norm in order to
prevent exploding gradients. Use None for no clipping.
exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
during training. Expect errors, silent or explicit, when `False`.
discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
from the prior. Training may be sped up by ignoring such less targeted
samples.
retrain_from_scratch_each_round: Whether to retrain the conditional density
estimator for the posterior from scratch each round.
show_train_summary: Whether to print the number of epochs and validation
loss and leakage after the training.
Returns:
Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
"""
kwargs = del_entries(locals(), entries=("self", "__class__"))
return super().train(**kwargs)
def __init__(
self,
prior,
density_estimator: Union[str, Callable] = "maf",
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[TensorboardSummaryWriter] = None,
show_progress_bars: bool = True,
**unused_args,
):
r"""Sequential Neural Likelihood [1].
[1] Sequential Neural Likelihood: Fast Likelihood-free Inference with
Autoregressive Flows_, Papamakarios et al., AISTATS 2019,
https://arxiv.org/abs/1805.07226
Args:
prior: A probability distribution that expresses prior knowledge about the
parameters, e.g. which ranges are meaningful for them. Any
object with `.log_prob()`and `.sample()` (for example, a PyTorch
distribution) can be used.
density_estimator: If it is a string, use a pre-configured network of the
provided type (one of nsf, maf, mdn, made). Alternatively, a function
that builds a custom neural network can be provided. The function will
be called with the first batch of simulations (theta, x), which can
thus be used for shape inference and potentially for z-scoring. It
needs to return a PyTorch `nn.Module` implementing the density
estimator. The density estimator needs to provide the methods
`.log_prob` and `.sample()`.
device: torch device on which to compute, e.g. gpu, cpu.
logging_level: Minimum severity of messages to log. One of the strings
INFO, WARNING, DEBUG, ERROR and CRITICAL.
summary_writer: A tensorboard `SummaryWriter` to control, among others, log
file location (default is `<current working directory>/logs`.)
show_progress_bars: Whether to show a progressbar during simulation and
sampling.
unused_args: Absorbs additional arguments. No entries will be used. If it
is not empty, we warn. In future versions, when the new interface of
0.14.0 is more mature, we will remove this argument.
"""
kwargs = del_entries(locals(),
entries=("self", "__class__", "unused_args"))
super().__init__(**kwargs, **unused_args)
def __init__(
self,
prior,
classifier: Union[str, Callable] = "resnet",
device: str = "cpu",
logging_level: Union[int, str] = "warning",
summary_writer: Optional[TensorboardSummaryWriter] = None,
show_progress_bars: bool = True,
**unused_args
):
r"""AALR[1], here known as SNRE_A.
[1] _Likelihood-free MCMC with Amortized Approximate Likelihood Ratios_, Hermans
et al., ICML 2020, https://arxiv.org/abs/1903.04057
Args:
prior: A probability distribution that expresses prior knowledge about the
parameters, e.g. which ranges are meaningful for them. Any
object with `.log_prob()`and `.sample()` (for example, a PyTorch
distribution) can be used.
classifier: Classifier trained to approximate likelihood ratios. If it is
a string, use a pre-configured network of the provided type (one of
linear, mlp, resnet). Alternatively, a function that builds a custom
neural network can be provided. The function will be called with the
first batch of simulations (theta, x), which can thus be used for shape
inference and potentially for z-scoring. It needs to return a PyTorch
`nn.Module` implementing the classifier.
device: torch device on which to compute, e.g. gpu, cpu.
logging_level: Minimum severity of messages to log. One of the strings
INFO, WARNING, DEBUG, ERROR and CRITICAL.
summary_writer: A tensorboard `SummaryWriter` to control, among others, log
file location (default is `<current working directory>/logs`.)
show_progress_bars: Whether to show a progressbar during simulation and
sampling.
unused_args: Absorbs additional arguments. No entries will be used. If it
is not empty, we warn. In future versions, when the new interface of
0.14.0 is more mature, we will remove this argument.
"""
kwargs = del_entries(locals(), entries=("self", "__class__", "unused_args"))
super().__init__(**kwargs, **unused_args)
def __init__(
self,
prior: Optional[Distribution] = None,
density_estimator: Union[str, Callable] = "maf",
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[TensorboardSummaryWriter] = None,
show_progress_bars: bool = True,
):
r"""SNPE-B [1]. CURRENTLY NOT IMPLEMENTED.
[1] _Flexible statistical inference for mechanistic models of neural dynamics_,
Lueckmann, Gonçalves et al., NeurIPS 2017, https://arxiv.org/abs/1711.01861.
See docstring of `PosteriorEstimator` class for all other arguments.
"""
raise NotImplementedError(
"SNPE-B is not yet implemented in the sbi package, see issue #199."
)
kwargs = del_entries(locals(), entries=("self", "__class__"))
super().__init__(**kwargs)
def __init__(
self,
prior: Optional[Distribution] = None,
classifier: Union[str, Callable] = "resnet",
device: str = "cpu",
logging_level: Union[int, str] = "warning",
summary_writer: Optional[TensorboardSummaryWriter] = None,
show_progress_bars: bool = True,
):
r"""AALR[1], here known as SNRE_A.
[1] _Likelihood-free MCMC with Amortized Approximate Likelihood Ratios_, Hermans
et al., ICML 2020, https://arxiv.org/abs/1903.04057
Args:
prior: A probability distribution that expresses prior knowledge about the
parameters, e.g. which ranges are meaningful for them. If `None`, the
prior must be passed to `.build_posterior()`.
classifier: Classifier trained to approximate likelihood ratios. If it is
a string, use a pre-configured network of the provided type (one of
linear, mlp, resnet). Alternatively, a function that builds a custom
neural network can be provided. The function will be called with the
first batch of simulations (theta, x), which can thus be used for shape
inference and potentially for z-scoring. It needs to return a PyTorch
`nn.Module` implementing the classifier.
device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
logging_level: Minimum severity of messages to log. One of the strings
INFO, WARNING, DEBUG, ERROR and CRITICAL.
summary_writer: A tensorboard `SummaryWriter` to control, among others, log
file location (default is `<current working directory>/logs`.)
show_progress_bars: Whether to show a progressbar during simulation and
sampling.
"""
kwargs = del_entries(locals(), entries=("self", "__class__"))
super().__init__(**kwargs)
def __call__(
self,
num_simulations: int,
proposal: Optional[Any] = None,
num_atoms: int = 10,
training_batch_size: int = 50,
learning_rate: float = 5e-4,
validation_fraction: float = 0.1,
stop_after_epochs: int = 20,
max_num_epochs: Optional[int] = None,
clip_max_norm: Optional[float] = 5.0,
calibration_kernel: Optional[Callable] = None,
exclude_invalid_x: bool = True,
discard_prior_samples: bool = False,
retrain_from_scratch_each_round: bool = False,
) -> DirectPosterior:
r"""Run SNPE.
Return posterior $p(\theta|x)$ after inference.
Args:
num_simulations: Number of simulator calls.
proposal: Distribution that the parameters $\theta$ are drawn from.
`proposal=None` uses the prior. Setting the proposal to a distribution
targeted on a specific observation, e.g. a posterior $p(\theta|x_o)$
obtained previously, can lead to less required simulations.
num_atoms: Number of atoms to use for classification.
training_batch_size: Training batch size.
learning_rate: Learning rate for Adam optimizer.
validation_fraction: The fraction of data to use for validation.
stop_after_epochs: The number of epochs to wait for improvement on the
validation set before terminating training.
max_num_epochs: Maximum number of epochs to run. If reached, we stop
training even when the validation loss is still decreasing. If None, we
train until validation loss increases (see also `stop_after_epochs`).
clip_max_norm: Value at which to clip the total gradient norm in order to
prevent exploding gradients. Use None for no clipping.
calibration_kernel: A function to calibrate the loss with respect to the
simulations `x`. See Lueckmann, Gonçalves et al., NeurIPS 2017.
exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
during training. Expect errors, silent or explicit, when `False`.
discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
from the prior. Training may be sped up by ignoring such less targeted
samples.
retrain_from_scratch_each_round: Whether to retrain the conditional density
estimator for the posterior from scratch each round.
Returns:
Posterior $p(\theta|x)$ that can be sampled and evaluated.
"""
# WARNING: sneaky trick ahead. We proxy the parent's `__call__` here,
# requiring the signature to have `num_atoms`, save it for use below, and
# continue. It's sneaky because we are using the object (self) as a namespace
# to pass arguments between functions, and that's implicit state management.
self._num_atoms = num_atoms
kwargs = del_entries(locals(),
entries=("self", "__class__", "num_atoms"))
if proposal is not None:
self.use_non_atomic_loss = (
isinstance(proposal.net._distribution, mdn)
and isinstance(self._posterior.net._distribution, mdn)
and (isinstance(self._prior, utils.BoxUniform)
or isinstance(self._prior, MultivariateNormal)))
algorithm = "non-atomic" if self.use_non_atomic_loss else "atomic"
print(f"Using SNPE-C with {algorithm} loss")
if self.use_non_atomic_loss:
# Take care of z-scoring, pre-compute and store prior terms.
self._set_state_for_mog_proposal()
return super().__call__(**kwargs)
def train(
self,
num_atoms: int = 10,
training_batch_size: int = 50,
learning_rate: float = 5e-4,
validation_fraction: float = 0.1,
stop_after_epochs: int = 20,
max_num_epochs: Optional[int] = None,
clip_max_norm: Optional[float] = 5.0,
calibration_kernel: Optional[Callable] = None,
exclude_invalid_x: bool = True,
resume_training: bool = False,
discard_prior_samples: bool = False,
use_combined_loss: bool = False,
retrain_from_scratch_each_round: bool = False,
show_train_summary: bool = False,
dataloader_kwargs: Optional[Dict] = None,
) -> DirectPosterior:
r"""
Return density estimator that approximates the distribution $p(\theta|x)$.
Args:
num_atoms: Number of atoms to use for classification.
training_batch_size: Training batch size.
learning_rate: Learning rate for Adam optimizer.
validation_fraction: The fraction of data to use for validation.
stop_after_epochs: The number of epochs to wait for improvement on the
validation set before terminating training.
max_num_epochs: Maximum number of epochs to run. If reached, we stop
training even when the validation loss is still decreasing. If None, we
train until validation loss increases (see also `stop_after_epochs`).
clip_max_norm: Value at which to clip the total gradient norm in order to
prevent exploding gradients. Use None for no clipping.
calibration_kernel: A function to calibrate the loss with respect to the
simulations `x`. See Lueckmann, Gonçalves et al., NeurIPS 2017.
exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
during training. Expect errors, silent or explicit, when `False`.
resume_training: Can be used in case training time is limited, e.g. on a
cluster. If `True`, the split between train and validation set, the
optimizer, the number of epochs, and the best validation log-prob will
be restored from the last time `.train()` was called.
discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
from the prior. Training may be sped up by ignoring such less targeted
samples.
use_combined_loss: Whether to train the neural net also on prior samples
using maximum likelihood in addition to training it on all samples using
atomic loss. The extra MLE loss helps prevent density leaking with
bounded priors.
retrain_from_scratch_each_round: Whether to retrain the conditional density
estimator for the posterior from scratch each round.
show_train_summary: Whether to print the number of epochs and validation
loss and leakage after the training.
dataloader_kwargs: Additional or updated kwargs to be passed to the training
and validation dataloaders (like, e.g., a collate_fn)
Returns:
Density estimator that approximates the distribution $p(\theta|x)$.
"""
# WARNING: sneaky trick ahead. We proxy the parent's `train` here,
# requiring the signature to have `num_atoms`, save it for use below, and
# continue. It's sneaky because we are using the object (self) as a namespace
# to pass arguments between functions, and that's implicit state management.
self._num_atoms = num_atoms
self._use_combined_loss = use_combined_loss
kwargs = del_entries(
locals(), entries=("self", "__class__", "num_atoms", "use_combined_loss")
)
self._round = max(self._data_round_index)
if self._round > 0:
# Set the proposal to the last proposal that was passed by the user. For
# atomic SNPE, it does not matter what the proposal is. For non-atomic
# SNPE, we only use the latest data that was passed, i.e. the one from the
# last proposal.
proposal = self._proposal_roundwise[-1]
if hasattr(proposal, "net"):
self.use_non_atomic_loss = (
isinstance(proposal.net._distribution, mdn)
and isinstance(self._neural_net._distribution, mdn)
and (
check_if_boxuniform(self._prior)[0]
or isinstance(self._prior, MultivariateNormal)
)
)
else:
self.use_non_atomic_loss = False
algorithm = "non-atomic" if self.use_non_atomic_loss else "atomic"
print(f"Using SNPE-C with {algorithm} loss")
if self.use_non_atomic_loss:
# Take care of z-scoring, pre-compute and store prior terms.
self._set_state_for_mog_proposal()
return super().train(**kwargs)
def __init__(
self,
simulator: Callable,
prior,
num_workers: int = 1,
simulation_batch_size: int = 1,
density_estimator: Union[str, Callable] = "maf",
mcmc_method: str = "slice_np",
mcmc_parameters: Optional[Dict[str, Any]] = None,
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[TensorboardSummaryWriter] = None,
show_progress_bars: bool = True,
show_round_summary: bool = False,
):
r"""Sequential Neural Likelihood [1].
[1] Sequential Neural Likelihood: Fast Likelihood-free Inference with
Autoregressive Flows_, Papamakarios et al., AISTATS 2019,
https://arxiv.org/abs/1805.07226
Args:
simulator: A function that takes parameters $\theta$ and maps them to
simulations, or observations, `x`, $\text{sim}(\theta)\to x$. Any
regular Python callable (i.e. function or class with `__call__` method)
can be used.
prior: A probability distribution that expresses prior knowledge about the
parameters, e.g. which ranges are meaningful for them. Any
object with `.log_prob()`and `.sample()` (for example, a PyTorch
distribution) can be used.
num_workers: Number of parallel workers to use for simulations.
simulation_batch_size: Number of parameter sets that the simulator
maps to data x at once. If None, we simulate all parameter sets at the
same time. If >= 1, the simulator has to process data of shape
(simulation_batch_size, parameter_dimension).
density_estimator: If it is a string, use a pre-configured network of the
provided type (one of nsf, maf, mdn, made). Alternatively, a function
that builds a custom neural network can be provided. The function will
be called with the first batch of simulations (theta, x), which can
thus be used for shape inference and potentially for z-scoring. It
needs to return a PyTorch `nn.Module` implementing the density
estimator. The density estimator needs to provide the methods
`.log_prob` and `.sample()`.
mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`, `hmc`, `nuts`.
Currently defaults to `slice_np` for a custom numpy implementation of
slice sampling; select `hmc`, `nuts` or `slice` for Pyro-based sampling.
mcmc_parameters: Dictionary overriding the default parameters for MCMC.
The following parameters are supported: `thin` to set the thinning
factor for the chain, `warmup_steps` to set the initial number of
samples to discard, `num_chains` for the number of chains, `init_strategy`
for the initialisation strategy for chains; `prior` will draw init
locations from prior, whereas `sir` will use Sequential-Importance-
Resampling using `init_strategy_num_candidates` to find init
locations.
device: torch device on which to compute, e.g. gpu, cpu.
logging_level: Minimum severity of messages to log. One of the strings
INFO, WARNING, DEBUG, ERROR and CRITICAL.
summary_writer: A tensorboard `SummaryWriter` to control, among others, log
file location (default is `<current working directory>/logs`.)
show_progress_bars: Whether to show a progressbar during simulation and
sampling.
show_round_summary: Whether to show the validation loss and leakage after
each round.
"""
kwargs = del_entries(locals(), entries=("self", "__class__"))
super().__init__(**kwargs)
def __init__(
self,
simulator: Callable,
prior,
num_workers: int = 1,
simulation_batch_size: int = 1,
density_estimator: Union[str, Callable] = "maf",
sample_with_mcmc: bool = False,
mcmc_method: str = "slice_np",
mcmc_parameters: Optional[Dict[str, Any]] = None,
use_combined_loss: bool = False,
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[TensorboardSummaryWriter] = None,
show_progress_bars: bool = True,
show_round_summary: bool = False,
):
r"""SNPE-C / APT [1].
[1] _Automatic Posterior Transformation for Likelihood-free Inference_,
Greenberg et al., ICML 2019, https://arxiv.org/abs/1905.07488.
This class implements two loss variants of SNPE-C: the non-atomic and the atomic
version. The atomic loss of SNPE-C can be used for any density estimator,
i.e. also for normalizing flows. However, it suffers from leakage issues. On
the other hand, the non-atomic loss can only be used only if the proposal
distribution is a mixture of Gaussians, the density estimator is a mixture of
Gaussians, and the prior is either Gaussian or Uniform. It does not suffer from
leakage issues. At the beginning of each round, we print whether the non-atomic
or the atomic version is used.
In this codebase, we will automatically switch to the non-atomic loss if the
following criteria are fulfilled:
- proposal has is a `DirectPosterior` with density_estimator `mdn`, as built
with `utils.sbi.posterior_nn()`.
- the density estimator is a `mdn`, as built with `utils.sbi.posterior_nn()`.
- `isinstance(prior, MultivariateNormal)` (from `torch.distributions`) or
`isinstance(prior, sbi.utils.BoxUniform)`
Note that custom implementations of any of these densities (or estimators) will
not trigger the non-atomic loss, and the algorithm will fall back onto using
the atomic loss.
Args:
simulator: A function that takes parameters $\theta$ and maps them to
simulations, or observations, `x`, $\mathrm{sim}(\theta)\to x$. Any
regular Python callable (i.e. function or class with `__call__` method)
can be used.
prior: A probability distribution that expresses prior knowledge about the
parameters, e.g. which ranges are meaningful for them. Any
object with `.log_prob()`and `.sample()` (for example, a PyTorch
distribution) can be used.
num_workers: Number of parallel workers to use for simulations.
simulation_batch_size: Number of parameter sets that the simulator
maps to data x at once. If None, we simulate all parameter sets at the
same time. If >= 1, the simulator has to process data of shape
(simulation_batch_size, parameter_dimension).
density_estimator: If it is a string, use a pre-configured network of the
provided type (one of nsf, maf, mdn, made). Alternatively, a function
that builds a custom neural network can be provided. The function will
be called with the first batch of simulations (theta, x), which can
thus be used for shape inference and potentially for z-scoring. It
needs to return a PyTorch `nn.Module` implementing the density
estimator. The density estimator needs to provide the methods
`.log_prob` and `.sample()`.
sample_with_mcmc: Whether to sample with MCMC. MCMC can be used to deal
with high leakage.
mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`,
`hmc`, `nuts`. Currently defaults to `slice_np` for a custom numpy
implementation of slice sampling; select `hmc`, `nuts` or `slice` for
Pyro-based sampling.
mcmc_parameters: Dictionary overriding the default parameters for MCMC.
The following parameters are supported: `thin` to set the thinning
factor for the chain, `warmup_steps` to set the initial number of
samples to discard, `num_chains` for the number of chains,
`init_strategy` for the initialisation strategy for chains; `prior` will
draw init locations from prior, whereas `sir` will use
Sequential-Importance-Resampling using `init_strategy_num_candidates`
to find init locations.
use_combined_loss: Whether to train the neural net also on prior samples
using maximum likelihood in addition to training it on all samples using
atomic loss. The extra MLE loss helps prevent density leaking with
bounded priors.
device: torch device on which to compute, e.g. gpu, cpu.
logging_level: Minimum severity of messages to log. One of the strings
INFO, WARNING, DEBUG, ERROR and CRITICAL.
summary_writer: A tensorboard `SummaryWriter` to control, among others, log
file location (default is `<current working directory>/logs`.)
show_progress_bars: Whether to show a progressbar during simulation and
sampling.
show_round_summary: Whether to show the validation loss and leakage after
each round.
"""
self._use_combined_loss = use_combined_loss
kwargs = del_entries(locals(),
entries=("self", "__class__",
"use_combined_loss"))
super().__init__(**kwargs)
def __init__(
self,
prior: Optional[Any] = None,
density_estimator: Union[str, Callable] = "mdn_snpe_a",
num_components: int = 10,
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[TensorboardSummaryWriter] = None,
show_progress_bars: bool = True,
**unused_args,
):
r"""SNPE-A [1].
[1] _Fast epsilon-free Inference of Simulation Models with Bayesian Conditional
Density Estimation_, Papamakarios et al., NeurIPS 2016,
https://arxiv.org/abs/1605.06376.
This class implements SNPE-A. SNPE-A trains across multiple rounds with a
maximum-likelihood-loss. This will make training converge to the proposal
posterior instead of the true posterior. To correct for this, SNPE-A applies a
post-hoc correction after training. This correction has to be performed
analytically. Thus, SNPE-A is limited to Gaussian distributions for all but the
last round. In the last round, SNPE-A can use a Mixture of Gaussians.
Args:
prior: A probability distribution that expresses prior knowledge about the
parameters, e.g. which ranges are meaningful for them. Any
object with `.log_prob()`and `.sample()` (for example, a PyTorch
distribution) can be used.
density_estimator: If it is a string (only "mdn_snpe_a" is valid), use a
pre-configured mixture of densities network. Alternatively, a function
that builds a custom neural network can be provided. The function will
be called with the first batch of simulations (theta, x), which can
thus be used for shape inference and potentially for z-scoring. It
needs to return a PyTorch `nn.Module` implementing the density
estimator. The density estimator needs to provide the methods
`.log_prob` and `.sample()`. Note that until the last round only a
single (multivariate) Gaussian component is used for training (see
Algorithm 1 in [1]). In the last round, this component is replicated
`num_components` times, its parameters are perturbed with a very small
noise, and then the last training round is done with the expanded
Gaussian mixture as estimator for the proposal posterior.
num_components: Number of components of the mixture of Gaussians in the
last round. This overrides the `num_components` value passed to
`posterior_nn()`.
device: torch device on which to compute, e.g. gpu, cpu.
logging_level: Minimum severity of messages to log. One of the strings
INFO, WARNING, DEBUG, ERROR and CRITICAL.
summary_writer: A tensorboard `SummaryWriter` to control, among others, log
file location (default is `<current working directory>/logs`.)
show_progress_bars: Whether to show a progressbar during training.
unused_args: Absorbs additional arguments. No entries will be used. If it
is not empty, we warn. In future versions, when the new interface of
0.14.0 is more mature, we will remove this argument.
"""
# Catch invalid inputs.
if not ((density_estimator == "mdn_snpe_a") or callable(density_estimator)):
raise TypeError(
"The `density_estimator` passed to SNPE_A needs to be a "
"callable or the string 'mdn_snpe_a'!"
)
# `num_components` will be used to replicate the Gaussian in the last round.
self._num_components = num_components
self._ran_final_round = False
# WARNING: sneaky trick ahead. We proxy the parent's `train` here,
# requiring the signature to have `num_atoms`, save it for use below, and
# continue. It's sneaky because we are using the object (self) as a namespace
# to pass arguments between functions, and that's implicit state management.
kwargs = utils.del_entries(
locals(),
entries=(
"self",
"__class__",
"unused_args",
"num_components",
),
)
super().__init__(**kwargs)
def train(
self,
final_round: bool = False,
training_batch_size: int = 50,
learning_rate: float = 5e-4,
validation_fraction: float = 0.1,
stop_after_epochs: int = 20,
max_num_epochs: Optional[int] = None,
clip_max_norm: Optional[float] = 5.0,
calibration_kernel: Optional[Callable] = None,
exclude_invalid_x: bool = True,
resume_training: bool = False,
retrain_from_scratch_each_round: bool = False,
show_train_summary: bool = False,
dataloader_kwargs: Optional[Dict] = None,
component_perturbation: float = 5e-3,
) -> DirectPosterior:
r"""
Return density estimator that approximates the proposal posterior $\tilde{p}(\theta|x)$.
[1] _Fast epsilon-free Inference of Simulation Models with Bayesian Conditional
Density Estimation_, Papamakarios et al., NeurIPS 2016,
https://arxiv.org/abs/1605.06376.
Training is performed with maximum likelihood on samples from the latest round,
which leads the algorithm to converge to the proposal posterior.
Args:
final_round: Whether we are in the last round of training or not. For all
but the last round, Algorithm 1 from [1] is executed. In last the
round, Algorithm 2 from [1] is executed once.
training_batch_size: Training batch size.
learning_rate: Learning rate for Adam optimizer.
validation_fraction: The fraction of data to use for validation.
stop_after_epochs: The number of epochs to wait for improvement on the
validation set before terminating training.
max_num_epochs: Maximum number of epochs to run. If reached, we stop
training even when the validation loss is still decreasing. If None, we
train until validation loss increases (see also `stop_after_epochs`).
clip_max_norm: Value at which to clip the total gradient norm in order to
prevent exploding gradients. Use None for no clipping.
calibration_kernel: A function to calibrate the loss with respect to the
simulations `x`. See Lueckmann, Gonçalves et al., NeurIPS 2017.
exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
during training. Expect errors, silent or explicit, when `False`.
resume_training: Can be used in case training time is limited, e.g. on a
cluster. If `True`, the split between train and validation set, the
optimizer, the number of epochs, and the best validation log-prob will
be restored from the last time `.train()` was called.
retrain_from_scratch_each_round: Whether to retrain the conditional density
estimator for the posterior from scratch each round. Not supported for
SNPE-A.
show_train_summary: Whether to print the number of epochs and validation
loss and leakage after the training.
dataloader_kwargs: Additional or updated kwargs to be passed to the training
and validation dataloaders (like, e.g., a collate_fn)
component_perturbation: The standard deviation applied to all weights and
biases when, in the last round, the Mixture of Gaussians is build from
a single Gaussian. This value can be problem-specific and also depends
on the number of mixture components.
Returns:
Density estimator that approximates the distribution $p(\theta|x)$.
"""
assert not retrain_from_scratch_each_round, """Retraining from scratch is not supported in SNPE-A yet. The reason for
this is that, if we reininitialized the density estimator, the z-scoring would
change, which would break the posthoc correction. This is a pure implementation
issue."""
kwargs = utils.del_entries(
locals(),
entries=("self", "__class__", "final_round", "component_perturbation"),
)
# SNPE-A always discards the prior samples.
kwargs["discard_prior_samples"] = True
self._round = max(self._data_round_index)
if final_round:
# If there is (will be) only one round, train with Algorithm 2 from [1].
if self._round == 0:
self._build_neural_net = partial(
self._build_neural_net, num_components=self._num_components
)
# Run Algorithm 2 from [1].
elif not self._ran_final_round:
# Now switch to the specified number of components. This method will
# only be used if `retrain_from_scratch_each_round=True`. Otherwise,
# the MDN will be built from replicating the single-component net for
# `num_component` times (via `_expand_mog()`).
self._build_neural_net = partial(
self._build_neural_net, num_components=self._num_components
)
# Extend the MDN to the originally desired number of components.
self._expand_mog(eps=component_perturbation)
else:
warnings.warn(
f"You have already run SNPE-A with `final_round=True`. Running it"
f"again with this setting will not allow computing the posthoc"
f"correction applied in SNPE-A. Thus, you will get an error when "
f"calling `.build_posterior()` after training.",
UserWarning,
)
else:
# Run Algorithm 1 from [1].
# Wrap the function that builds the MDN such that we can make
# sure that there is only one component when running.
self._build_neural_net = partial(self._build_neural_net, num_components=1)
if final_round:
self._ran_final_round = True
return super().train(**kwargs)
def __call__(
self,
num_rounds: int,
num_simulations_per_round: OneOrMore[int],
x_o: Optional[Tensor] = None,
num_atoms: int = 10,
training_batch_size: int = 50,
learning_rate: float = 5e-4,
validation_fraction: float = 0.1,
stop_after_epochs: int = 20,
max_num_epochs: Optional[int] = None,
clip_max_norm: Optional[float] = 5.0,
calibration_kernel: Optional[Callable] = None,
exclude_invalid_x: bool = True,
discard_prior_samples: bool = False,
retrain_from_scratch_each_round: bool = False,
) -> NeuralPosterior:
r"""Run SNPE.
Return posterior $p(\theta|x)$ after inference (possibly over several rounds).
Args:
num_rounds: Number of rounds to run. Each round consists of a simulation and
training phase. `num_rounds=1` leads to a posterior $p(\theta|x)$ valid
for _any_ $x$ (amortized), but requires many simulations.
Alternatively, with `num_rounds>1` the inference returns a posterior
$p(\theta|x_o)$ focused on a specific observation `x_o`, potentially
requiring less simulations.
num_simulations_per_round: Number of simulator calls per round.
x_o: An observation that is only required when doing inference
over multiple rounds. After the first round, `x_o` is used to guide the
sampling so that the simulator is run with parameters that are likely
for that `x_o`, i.e. they are sampled from the posterior obtained in the
previous round $p(\theta|x_o)$.
num_atoms: Number of atoms to use for classification.
training_batch_size: Training batch size.
learning_rate: Learning rate for Adam optimizer.
validation_fraction: The fraction of data to use for validation.
stop_after_epochs: The number of epochs to wait for improvement on the
validation set before terminating training.
max_num_epochs: Maximum number of epochs to run. If reached, we stop
training even when the validation loss is still decreasing. If None, we
train until validation loss increases (see also `stop_after_epochs`).
clip_max_norm: Value at which to clip the total gradient norm in order to
prevent exploding gradients. Use None for no clipping.
calibration_kernel: A function to calibrate the loss with respect to the
simulations `x`. See Lueckmann, Gonçalves et al., NeurIPS 2017.
exclude_invalid_x: Whether to exclude simulation outputs `x=NaN` or `x=±∞`
during training. Expect errors, silent or explicit, when `False`.
discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
from the prior. Training may be sped up by ignoring such less targeted
samples.
retrain_from_scratch_each_round: Whether to retrain the conditional density
estimator for the posterior from scratch each round.
Returns:
Posterior $p(\theta|x)$ that can be sampled and evaluated.
"""
# WARNING: sneaky trick ahead. We proxy the parent's `__call__` here,
# requiring the signature to have `num_atoms`, save it for use below, and
# continue. It's sneaky because we are using the object (self) as a namespace
# to pass arguments between functions, and that's implicit state management.
self._num_atoms = num_atoms
kwargs = del_entries(locals(),
entries=("self", "__class__", "num_atoms"))
return super().__call__(**kwargs)
