def test_process_device(device_input: str, device_target: Optional[str]) -> None:
device_output = process_device(device_input)
assert device_output == device_target, (
f"Failure when processing device '{device_input}': "
f"result should have been '{device_target}' and is "
f"instead '{device_output}'"
)
def test_training_and_mcmc_on_device(method, model, device):
"""Test training on devices.
This test does not check training speeds.
"""
device = process_device(device)
num_dim = 2
num_samples = 10
num_simulations = 500
max_num_epochs = 5
x_o = zeros(1, num_dim)
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * eye(num_dim)
prior_mean = zeros(num_dim)
prior_cov = eye(num_dim)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
if method == SNPE:
kwargs = dict(density_estimator=utils.posterior_nn(model=model), )
mcmc_kwargs = dict(
sample_with_mcmc=True,
mcmc_method="slice_np",
)
elif method == SNLE:
kwargs = dict(density_estimator=utils.likelihood_nn(model=model), )
mcmc_kwargs = dict(mcmc_method="slice")
elif method == SNRE:
kwargs = dict(classifier=utils.classifier_nn(model=model), )
mcmc_kwargs = dict(mcmc_method="slice_np_vectorized", )
else:
raise ValueError()
inferer = method(prior, show_progress_bars=False, device=device, **kwargs)
proposals = [prior]
# Test for two rounds.
for r in range(2):
theta, x, = simulate_for_sbi(simulator,
proposal=prior,
num_simulations=num_simulations)
_ = inferer.append_simulations(theta,
x).train(training_batch_size=100,
max_num_epochs=max_num_epochs)
posterior = inferer.build_posterior(**mcmc_kwargs).set_default_x(x_o)
proposals.append(posterior)
proposals[-1].sample(sample_shape=(num_samples, ), x=x_o, **mcmc_kwargs)
def __init__(
self,
potential_fn: Callable,
theta_transform: Optional[TorchTransform] = None,
device: Optional[str] = None,
x_shape: Optional[torch.Size] = None,
):
"""
Args:
potential_fn: The potential function from which to draw samples.
theta_transform: Transformation that will be applied during sampling.
Allows to perform, e.g. MCMC in unconstrained space.
device: Training device, e.g., "cpu", "cuda" or "cuda:0". If None,
`potential_fn.device` is used.
"""
# Ensure device string.
self._device = process_device(
potential_fn.device if device is None else device)
self.potential_fn = potential_fn
if theta_transform is None:
self.theta_transform = torch_tf.IndependentTransform(
torch_tf.identity_transform, reinterpreted_batch_ndims=1)
else:
self.theta_transform = theta_transform
self._map = None
self._purpose = ""
self._x_shape = x_shape
# If the sampler interface (#573) is used, the user might have passed `x_o`
# already to the potential function builder. If so, this `x_o` will be used
# as default x.
self._x = self.potential_fn.return_x_o()
def __init__(
self,
prior: Optional[Distribution] = None,
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[SummaryWriter] = None,
show_progress_bars: bool = True,
):
r"""Base class for inference methods.
Args:
prior: A probability distribution that expresses prior knowledge about the
parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
distribution, see FAQ for details on how to use custom distributions.
device: torch device on which to train the neural net and on which to
perform all posterior operations, e.g. gpu or cpu.
logging_level: Minimum severity of messages to log. One of the strings
"INFO", "WARNING", "DEBUG", "ERROR" and "CRITICAL".
summary_writer: A `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.
"""
self._device = process_device(device)
check_prior(prior)
check_if_prior_on_device(self._device, prior)
self._prior = prior
self._posterior = None
self._neural_net = None
self._x_shape = None
self._show_progress_bars = show_progress_bars
# Initialize roundwise (theta, x, prior_masks) for storage of parameters,
# simulations and masks indicating if simulations came from prior.
self._theta_roundwise, self._x_roundwise, self._prior_masks = [], [], []
self._model_bank = []
# Initialize list that indicates the round from which simulations were drawn.
self._data_round_index = []
self._round = 0
self._val_log_prob = float("-Inf")
# XXX We could instantiate here the Posterior for all children. Two problems:
# 1. We must dispatch to right PotentialProvider for mcmc based on name
# 2. `method_family` cannot be resolved only from `self.__class__.__name__`,
# since SRE, AALR demand different handling but are both in SRE class.
self._summary_writer = (self._default_summary_writer()
if summary_writer is None else summary_writer)
# Logging during training (by SummaryWriter).
self._summary = dict(
median_observation_distances=[],
epochs=[],
best_validation_log_probs=[],
validation_log_probs=[],
train_log_probs=[],
epoch_durations_sec=[],
)
def __init__(
self,
prior: Optional[Any] = None,
device: str = "cpu",
logging_level: Union[int, str] = "WARNING",
summary_writer: Optional[SummaryWriter] = None,
show_progress_bars: bool = True,
**unused_args,
):
r"""
Base class for inference methods.
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.
device: torch device on which to train the neural net and on which to
perform all posterior operations, e.g. gpu or cpu.
logging_level: Minimum severity of messages to log. One of the strings
"INFO", "WARNING", "DEBUG", "ERROR" and "CRITICAL".
summary_writer: A `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.
"""
self._device = process_device(device)
if unused_args:
warn(
f"You passed some keyword arguments that will not be used. "
f"Specifically, the unused arguments are: {list(unused_args.keys())}. "
f"These arguments might have been supported in sbi "
f"versions <0.14.0. Since 0.14.0, the API was changed. Please consult "
f"the corresponding pull request on github: "
f"https://github.com/mackelab/sbi/pull/378 and tutorials: "
f"https://www.mackelab.org/sbi/tutorial/02_flexible_interface/ for "
f"further information.",
)
self._prior = prior
self._posterior = None
self._neural_net = None
self._x_shape = None
self._show_progress_bars = show_progress_bars
# Initialize roundwise (theta, x, prior_masks) for storage of parameters,
# simulations and masks indicating if simulations came from prior.
self._theta_roundwise, self._x_roundwise, self._prior_masks = [], [], []
self._model_bank = []
# Initialize list that indicates the round from which simulations were drawn.
self._data_round_index = []
self._round = 0
# XXX We could instantiate here the Posterior for all children. Two problems:
# 1. We must dispatch to right PotentialProvider for mcmc based on name
# 2. `method_family` cannot be resolved only from `self.__class__.__name__`,
# since SRE, AALR demand different handling but are both in SRE class.
self._summary_writer = (
self._default_summary_writer() if summary_writer is None else summary_writer
)
# Logging during training (by SummaryWriter).
self._summary = dict(
median_observation_distances=[],
epochs=[],
best_validation_log_probs=[],
validation_log_probs=[],
train_log_probs=[],
epoch_durations_sec=[]
)
def test_process_device(device: str):
process_device(device)
def test_training_and_mcmc_on_device(
method,
model,
data_device,
mcmc_method,
training_device,
prior_device,
prior_type="gaussian",
):
"""Test training on devices.
This test does not check training speeds.
"""
num_dim = 2
num_samples = 10
num_simulations = 100
max_num_epochs = 5
x_o = zeros(1, num_dim).to(data_device)
likelihood_shift = -1.0 * ones(num_dim).to(prior_device)
likelihood_cov = 0.3 * eye(num_dim).to(prior_device)
if prior_type == "gaussian":
prior_mean = zeros(num_dim).to(prior_device)
prior_cov = eye(num_dim).to(prior_device)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
else:
prior = BoxUniform(
low=-2 * torch.ones(num_dim),
high=2 * torch.ones(num_dim),
device=prior_device,
)
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
training_device = process_device(training_device)
if method in [SNPE_A, SNPE_C]:
kwargs = dict(
density_estimator=utils.posterior_nn(model=model, num_transforms=2)
)
elif method == SNLE:
kwargs = dict(
density_estimator=utils.likelihood_nn(model=model, num_transforms=2)
)
elif method in (SNRE_A, SNRE_B):
kwargs = dict(classifier=utils.classifier_nn(model=model))
else:
raise ValueError()
inferer = method(show_progress_bars=False, device=training_device, **kwargs)
proposals = [prior]
# Test for two rounds.
for _ in range(2):
theta, x = simulate_for_sbi(simulator, proposals[-1], num_simulations)
theta, x = theta.to(data_device), x.to(data_device)
estimator = inferer.append_simulations(theta, x).train(
training_batch_size=100, max_num_epochs=max_num_epochs
)
if method == SNLE:
potential_fn, theta_transform = likelihood_estimator_based_potential(
estimator, prior, x_o
)
elif method == SNPE_A or method == SNPE_C:
potential_fn, theta_transform = posterior_estimator_based_potential(
estimator, prior, x_o
)
elif method == SNRE_A or method == SNRE_B:
potential_fn, theta_transform = ratio_estimator_based_potential(
estimator, prior, x_o
)
else:
raise ValueError
if mcmc_method == "rejection":
posterior = RejectionPosterior(
proposal=prior,
potential_fn=potential_fn,
device=training_device,
)
elif mcmc_method == "direct":
posterior = DirectPosterior(
posterior_estimator=estimator, prior=prior
).set_default_x(x_o)
else:
posterior = MCMCPosterior(
potential_fn=potential_fn,
theta_transform=theta_transform,
proposal=prior,
method=mcmc_method,
device=training_device,
)
proposals.append(posterior)
# Check for default device for inference object
weights_device = next(inferer._neural_net.parameters()).device
assert torch.device(training_device) == weights_device
samples = proposals[-1].sample(sample_shape=(num_samples,))
proposals[-1].potential(samples)
def test_training_and_mcmc_on_device(method, model, data_device, training_device):
"""Test training on devices.
This test does not check training speeds.
"""
training_device = process_device(training_device)
num_dim = 2
num_samples = 10
num_simulations = 500
max_num_epochs = 5
x_o = zeros(1, num_dim).to(data_device)
likelihood_shift = -1.0 * ones(num_dim)
likelihood_cov = 0.3 * eye(num_dim)
prior_mean = zeros(num_dim)
prior_cov = eye(num_dim)
prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov)
def simulator(theta):
return linear_gaussian(theta, likelihood_shift, likelihood_cov)
if method in [SNPE_A, SNPE_C]:
kwargs = dict(
density_estimator=utils.posterior_nn(model=model),
)
mcmc_kwargs = (
dict(
sample_with="mcmc",
mcmc_method="slice_np",
)
if method == SNPE_C
else {}
)
elif method == SNLE:
kwargs = dict(
density_estimator=utils.likelihood_nn(model=model),
)
mcmc_kwargs = dict(sample_with="mcmc", mcmc_method="slice")
elif method in (SNRE_A, SNRE_B):
kwargs = dict(
classifier=utils.classifier_nn(model=model),
)
mcmc_kwargs = dict(
sample_with="mcmc",
mcmc_method="slice_np_vectorized",
)
else:
raise ValueError()
inferer = method(prior, show_progress_bars=False, device=training_device, **kwargs)
proposals = [prior]
# Test for two rounds.
for _ in range(2):
theta, x = simulate_for_sbi(simulator, prior, num_simulations)
theta, x = theta.to(data_device), x.to(data_device)
_ = inferer.append_simulations(theta, x).train(
training_batch_size=100, max_num_epochs=max_num_epochs
)
posterior = inferer.build_posterior(**mcmc_kwargs).set_default_x(x_o)
proposals.append(posterior)
# Check for default device for inference object
weights_device = next(inferer._neural_net.parameters()).device
assert torch.device(training_device) == weights_device
samples = proposals[-1].sample(sample_shape=(num_samples,), x=x_o, **mcmc_kwargs)
proposals[-1].log_prob(samples)