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,
) -> RatioBasedPosterior:
r"""
Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
Args:
num_atoms: Number of atoms to use for classification.
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:
Classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
"""
max_num_epochs = 2**31 - 1 if max_num_epochs is None else max_num_epochs
# Load data from most recent round.
self._round = max(self._data_round_index)
theta, x, _ = self.get_simulations(self._round, exclude_invalid_x,
False)
# Starting index for the training set (1 = discard round-0 samples).
start_idx = int(discard_prior_samples and self._round > 0)
theta, x, _ = self.get_simulations(start_idx, exclude_invalid_x)
# Get total number of training examples.
num_examples = len(theta)
# Select random train and validation splits from (theta, x) pairs.
permuted_indices = torch.randperm(num_examples)
num_training_examples = int((1 - validation_fraction) * num_examples)
num_validation_examples = num_examples - num_training_examples
train_indices, val_indices = (
permuted_indices[:num_training_examples],
permuted_indices[num_training_examples:],
)
clipped_batch_size = min(training_batch_size, num_validation_examples)
num_atoms = clamp_and_warn("num_atoms",
num_atoms,
min_val=2,
max_val=clipped_batch_size)
# Dataset is shared for training and validation loaders.
dataset = data.TensorDataset(theta, x)
# Create neural net and validation loaders using a subset sampler.
train_loader = data.DataLoader(
dataset,
batch_size=clipped_batch_size,
drop_last=True,
sampler=SubsetRandomSampler(train_indices),
)
val_loader = data.DataLoader(
dataset,
batch_size=clipped_batch_size,
shuffle=False,
drop_last=False,
sampler=SubsetRandomSampler(val_indices),
)
# First round or if retraining from scratch:
# Call the `self._build_neural_net` with the rounds' thetas and xs as
# arguments, which will build the neural network
# This is passed into NeuralPosterior, to create a neural posterior which
# can `sample()` and `log_prob()`. The network is accessible via `.net`.
if self._neural_net is None or retrain_from_scratch_each_round:
self._neural_net = self._build_neural_net(theta[train_indices],
x[train_indices])
self._x_shape = x_shape_from_simulation(x)
assert len(self._x_shape) < 3, (
"For now, SNRE cannot handle multi-dimensional simulator output, see "
"issue #360.")
self._neural_net.to(self._device)
optimizer = optim.Adam(
list(self._neural_net.parameters()),
lr=learning_rate,
)
epoch, self._val_log_prob = 0, float("-Inf")
while epoch <= max_num_epochs and not self._converged(
epoch, stop_after_epochs):
# Train for a single epoch.
self._neural_net.train()
for batch in train_loader:
optimizer.zero_grad()
theta_batch, x_batch = (
batch[0].to(self._device),
batch[1].to(self._device),
)
loss = self._loss(theta_batch, x_batch, num_atoms)
loss.backward()
if clip_max_norm is not None:
clip_grad_norm_(
self._neural_net.parameters(),
max_norm=clip_max_norm,
)
optimizer.step()
epoch += 1
# Calculate validation performance.
self._neural_net.eval()
log_prob_sum = 0
with torch.no_grad():
for batch in val_loader:
theta_batch, x_batch = (
batch[0].to(self._device),
batch[1].to(self._device),
)
log_prob = self._loss(theta_batch, x_batch, num_atoms)
log_prob_sum -= log_prob.sum().item()
self._val_log_prob = log_prob_sum / num_validation_examples
# Log validation log prob for every epoch.
self._summary["validation_log_probs"].append(
self._val_log_prob)
self._maybe_show_progress(self._show_progress_bars, epoch)
self._report_convergence_at_end(epoch, stop_after_epochs,
max_num_epochs)
# Update summary.
self._summary["epochs"].append(epoch)
self._summary["best_validation_log_probs"].append(
self._best_val_log_prob)
# Update TensorBoard and summary dict.
self._summarize(
round_=self._round,
x_o=None,
theta_bank=theta,
x_bank=x,
)
# Update description for progress bar.
if show_train_summary:
print(self._describe_round(self._round, self._summary))
return deepcopy(self._neural_net)
def _log_prob_proposal_posterior_atomic(
self, theta: Tensor, x: Tensor, masks: Tensor
):
"""
Return log probability of the proposal posterior for atomic proposals.
We have two main options when evaluating the proposal posterior.
(1) Generate atoms from the proposal prior.
(2) Generate atoms from a more targeted distribution, such as the most
recent posterior.
If we choose the latter, it is likely beneficial not to do this in the first
round, since we would be sampling from a randomly-initialized neural density
estimator.
Args:
theta: Batch of parameters θ.
x: Batch of data.
masks: Mask that is True for prior samples in the batch in order to train
them with prior loss.
Returns:
Log-probability of the proposal posterior.
"""
batch_size = theta.shape[0]
num_atoms = clamp_and_warn(
"num_atoms", self._num_atoms, min_val=2, max_val=batch_size
)
# Each set of parameter atoms is evaluated using the same x,
# so we repeat rows of the data x, e.g. [1, 2] -> [1, 1, 2, 2]
repeated_x = repeat_rows(x, num_atoms)
# To generate the full set of atoms for a given item in the batch,
# we sample without replacement num_atoms - 1 times from the rest
# of the theta in the batch.
probs = ones(batch_size, batch_size) * (1 - eye(batch_size)) / (batch_size - 1)
choices = torch.multinomial(probs, num_samples=num_atoms - 1, replacement=False)
contrasting_theta = theta[choices]
# We can now create our sets of atoms from the contrasting parameter sets
# we have generated.
atomic_theta = torch.cat((theta[:, None, :], contrasting_theta), dim=1).reshape(
batch_size * num_atoms, -1
)
# Evaluate large batch giving (batch_size * num_atoms) log prob posterior evals.
log_prob_posterior = self._neural_net.log_prob(atomic_theta, repeated_x)
self._assert_all_finite(log_prob_posterior, "posterior eval")
log_prob_posterior = log_prob_posterior.reshape(batch_size, num_atoms)
# Get (batch_size * num_atoms) log prob prior evals.
log_prob_prior = self._prior.log_prob(atomic_theta)
log_prob_prior = log_prob_prior.reshape(batch_size, num_atoms)
self._assert_all_finite(log_prob_prior, "prior eval")
# Compute unnormalized proposal posterior.
unnormalized_log_prob = log_prob_posterior - log_prob_prior
# Normalize proposal posterior across discrete set of atoms.
log_prob_proposal_posterior = unnormalized_log_prob[:, 0] - torch.logsumexp(
unnormalized_log_prob, dim=-1
)
self._assert_all_finite(log_prob_proposal_posterior, "proposal posterior eval")
# XXX This evaluates the posterior on _all_ prior samples
if self._use_combined_loss:
log_prob_posterior_non_atomic = self._neural_net.log_prob(theta, x)
masks = masks.reshape(-1)
log_prob_proposal_posterior = (
masks * log_prob_posterior_non_atomic + log_prob_proposal_posterior
)
return log_prob_proposal_posterior
def _train(
self,
round_: int,
num_atoms: int,
training_batch_size: int,
learning_rate: float,
validation_fraction: float,
stop_after_epochs: int,
max_num_epochs: int,
clip_max_norm: Optional[float],
exclude_invalid_x: bool,
discard_prior_samples: bool,
) -> None:
r"""
Trains the neural classifier.
Update the classifier weights by maximizing a Bernoulli likelihood which
distinguishes between jointly distributed $(\theta, x)$ pairs and randomly
chosen $(\theta, x)$ pairs.
Uses performance on a held-out validation set as a terminating condition (early
stopping).
"""
# Starting index for the training set (1 = discard round-0 samples).
start_idx = int(discard_prior_samples and round_ > 0)
theta, x, _ = self._get_from_data_bank(start_idx, exclude_invalid_x)
# Get total number of training examples.
num_examples = len(theta)
# Select random train and validation splits from (theta, x) pairs.
permuted_indices = torch.randperm(num_examples)
num_training_examples = int((1 - validation_fraction) * num_examples)
num_validation_examples = num_examples - num_training_examples
train_indices, val_indices = (
permuted_indices[:num_training_examples],
permuted_indices[num_training_examples:],
)
clipped_batch_size = min(training_batch_size, num_validation_examples)
# num_atoms = theta.shape[0]
clamp_and_warn("num_atoms",
num_atoms,
min_val=2,
max_val=clipped_batch_size)
# Dataset is shared for training and validation loaders.
dataset = data.TensorDataset(theta, x)
# Create neural net and validation loaders using a subset sampler.
train_loader = data.DataLoader(
dataset,
batch_size=clipped_batch_size,
drop_last=True,
sampler=SubsetRandomSampler(train_indices),
)
val_loader = data.DataLoader(
dataset,
batch_size=clipped_batch_size,
shuffle=False,
drop_last=False,
sampler=SubsetRandomSampler(val_indices),
)
optimizer = optim.Adam(
list(self._posterior.net.parameters()),
lr=learning_rate,
)
epoch, self._val_log_prob = 0, float("-Inf")
while epoch <= max_num_epochs and not self._converged(
epoch, stop_after_epochs):
# Train for a single epoch.
self._posterior.net.train()
for batch in train_loader:
optimizer.zero_grad()
theta_batch, x_batch = (
batch[0].to(self._device),
batch[1].to(self._device),
)
loss = self._loss(theta_batch, x_batch, num_atoms)
loss.backward()
if clip_max_norm is not None:
clip_grad_norm_(
self._posterior.net.parameters(),
max_norm=clip_max_norm,
)
optimizer.step()
epoch += 1
# Calculate validation performance.
self._posterior.net.eval()
log_prob_sum = 0
with torch.no_grad():
for batch in val_loader:
theta_batch, x_batch = (
batch[0].to(self._device),
batch[1].to(self._device),
)
log_prob = self._loss(theta_batch, x_batch, num_atoms)
log_prob_sum -= log_prob.sum().item()
self._val_log_prob = log_prob_sum / num_validation_examples
self._maybe_show_progress(self._show_progress_bars, epoch)
self._report_convergence_at_end(epoch, stop_after_epochs,
max_num_epochs)
# Update summary.
self._summary["epochs"].append(epoch)
self._summary["best_validation_log_probs"].append(
self._best_val_log_prob)
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,
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,
) -> RatioBasedPosterior:
r"""
Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
Args:
num_atoms: Number of atoms to use for classification.
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.
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)$.
"""
max_num_epochs = 2 ** 31 - 1 if max_num_epochs is None else max_num_epochs
# Starting index for the training set (1 = discard round-0 samples).
start_idx = int(discard_prior_samples and self._round > 0)
# Load data from most recent round.
self._round = max(self._data_round_index)
theta, x, _ = self.get_simulations(
start_idx, exclude_invalid_x, warn_on_invalid=True
)
# Dataset is shared for training and validation loaders.
dataset = data.TensorDataset(theta, x)
train_loader, val_loader = self.get_dataloaders(
dataset,
training_batch_size,
validation_fraction,
resume_training,
dataloader_kwargs=dataloader_kwargs,
)
clipped_batch_size = min(training_batch_size, len(val_loader))
num_atoms = clamp_and_warn(
"num_atoms", num_atoms, min_val=2, max_val=clipped_batch_size
)
# First round or if retraining from scratch:
# Call the `self._build_neural_net` with the rounds' thetas and xs as
# arguments, which will build the neural network
# This is passed into NeuralPosterior, to create a neural posterior which
# can `sample()` and `log_prob()`. The network is accessible via `.net`.
if self._neural_net is None or retrain_from_scratch_each_round:
self._neural_net = self._build_neural_net(
theta[self.train_indices], x[self.train_indices]
)
self._x_shape = x_shape_from_simulation(x)
self._neural_net.to(self._device)
if not resume_training:
self.optimizer = optim.Adam(
list(self._neural_net.parameters()), lr=learning_rate,
)
self.epoch, self._val_log_prob = 0, float("-Inf")
while self.epoch <= max_num_epochs and not self._converged(
self.epoch, stop_after_epochs
):
# Train for a single epoch.
self._neural_net.train()
for batch in train_loader:
self.optimizer.zero_grad()
theta_batch, x_batch = (
batch[0].to(self._device),
batch[1].to(self._device),
)
loss = self._loss(theta_batch, x_batch, num_atoms)
loss.backward()
if clip_max_norm is not None:
clip_grad_norm_(
self._neural_net.parameters(), max_norm=clip_max_norm,
)
self.optimizer.step()
self.epoch += 1
# Calculate validation performance.
self._neural_net.eval()
loss_sum = 0
with torch.no_grad():
for batch in val_loader:
theta_batch, x_batch = (
batch[0].to(self._device),
batch[1].to(self._device),
)
loss = self._loss(theta_batch, x_batch, num_atoms)
loss_sum -= loss.sum().item()
# Take mean over all validation samples.
self._val_log_prob = loss_sum / (
len(val_loader) * val_loader.batch_size
)
# Log validation log prob for every epoch.
self._summary["validation_log_probs"].append(self._val_log_prob)
self._maybe_show_progress(self._show_progress_bars, self.epoch)
self._report_convergence_at_end(self.epoch, stop_after_epochs, max_num_epochs)
# Update summary.
self._summary["epochs"].append(self.epoch)
self._summary["best_validation_log_probs"].append(self._best_val_log_prob)
# Update TensorBoard and summary dict.
self._summarize(
round_=self._round, x_o=None, theta_bank=theta, x_bank=x,
)
# Update description for progress bar.
if show_train_summary:
print(self._describe_round(self._round, self._summary))
return deepcopy(self._neural_net)