bayesflow.trainers module

class bayesflow.trainers.Trainer(amortizer, generative_model=None, configurator=None, checkpoint_path=None, max_to_keep=3, default_lr=0.0005, skip_checks=False, memory=False, **kwargs)

Bases: object

This class connects a generative model (or, already simulated data from a model) with a configurator and a neural inference architecture for amortized inference (amortizer). A Trainer instance is responsible for optimizing the amortizer via various forms of simulation-based training.

At the very minimum, the trainer must be initialized with an amortizer instance, which is capable of processing the (configured) outputs of a generative model. A configurator will then process the outputs of the generative model and convert them into suitable inputs for the amortizer. Users can choose from a palette of default configurators or create their own configurators, essentially building a modularized pipeline GenerativeModel -> Configurator -> Amortizer. Most complex models will require custom configurators.

Notes

Currently, the trainer supports the following simulation-based training regimes, based on efficiency considerations:

• Online training

  >>> trainer.train_online(epochs, iterations_per_epoch, batch_size, **kwargs)
  

  This training regime is optimal for fast generative models which can efficiently simulated data on-the-fly. In order for this training regime to be efficient, on-the-fly batch simulations should not take longer than 2-3 seconds.

• Experience replay training

  >>> trainer.train_experience_replay(epochs, iterations_per_epoch, batch_size, **kwargs)
  

  This training regime is also good for fast generative models capable of efficiently simulating data on-the-fly. Compare to pure online training, this training will keep an experience replay buffer from which simulations are randomly sampled, so the networks will likely see some simulations multiple times.

• Round-based training

  >>> trainer.train_rounds(rounds, sim_per_round, epochs, batch_size, **kwargs)
  

  This training regime is optimal for slow, but still reasonably performant generative models. In order for this training regime to be efficient, on-the-fly batch simulations should not take longer than 2-3 minutes.

  Note

  overfitting presents a danger when using small numbers of simulated data sets, so it is recommended to use some amount of regularization for the neural amortizer(s).

• Offline training

  >>> trainer.train_offline(simulations_dict, epochs, batch_size, **kwargs)
  

  This training regime is optimal for very slow, external simulators, which take several minutes for a single simulation. It assumes that all training data has been already simulated and stored on disk.

  Warning

  Overfitting presents a danger when using a small simulated data set, so it is recommended to use some amount of regularization for the neural amortizer(s).

  Note

  For extremely slow simulators (i.e., more than an hour of a single simulation), the BayesFlow framework might not be the ideal choice and should probably be considered in combination with a black-box surrogate optimization method, such as Bayesian optimization.

__init__(amortizer, generative_model=None, configurator=None, checkpoint_path=None, max_to_keep=3, default_lr=0.0005, skip_checks=False, memory=False, **kwargs)

Creates a trainer which will use a generative model (or data simulated from it) to optimize a neural architecture (amortizer) for amortized posterior inference, likelihood inference, or both.

Parameters:

amortizerbayesflow.amortizers.Amortizer

The neural architecture to be optimized.

generative_modelbayesflow.forward_inference.GenerativeModel

A generative model returning a dictionary with randomly sampled parameters, data, and optional context

configuratorcallable or None, optional, default: None

A callable object transforming and combining the outputs of the generative model into inputs for a BayesFlow amortizer.

checkpoint_pathstring or None, optional, default: None

Optional file path for storing the trained amortizer, loss history and optional memory.

max_to_keepint, optional, default: 3

Number of checkpoints and loss history snapshots to keep.

default_lrfloat, optional, default: 0.0005

The default learning rate to use for default optimizers.

skip_checksbool, optional, default: False

If True, do not perform consistency checks, i.e., simulator runs and passed through nets

memorybool or bayesflow.SimulationMemory, optional, default: False

If True, store a pre-defined amount of simulations for later use (validation, etc.). If SimulationMemory instance provided, stores a reference to the instance. Otherwise the corresponding attribute will be set to None.

diagnose_latent2d(inputs=None, **kwargs)

Performs visual pre-inference diagnostics of latent space on either provided validation data (new simulations) or internal simulation memory. If inputs is not None, then diagnostics will be performed on the inputs, regardless whether the simulation_memory of the trainer is empty or not. If inputs is None, then the trainer will try to access is memory or raise a ConfigurationError.

Parameters:

inputsNone, list, or dict, optional, default: None

The optional inputs to use

Returns:

figplt.Figure

The figure object which can be readily saved to disk using fig.savefig().

Other Parameters:

conf_args

optional keyword arguments passed to the configurator

net_args

optional keyword arguments passed to the amortizer

plot_args

optional keyword arguments passed to plot_latent_space_2d

diagnose_sbc_histograms(inputs=None, n_samples=None, **kwargs)

Performs visual pre-inference diagnostics via simulation-based calibration (SBC) (new simulations) or internal simulation memory. If inputs is not None, then diagnostics will be performed on the inputs, regardless whether the simulation_memory of the trainer is empty or not. If inputs is None, then the trainer will try to access is memory or raise a ConfigurationError.

Parameters:

inputsNone, list or dict, optional, default: None

The optional inputs to use

n_samplesint or None, optional, default: None

The number of posterior samples to draw for each simulated data set. If None, the number will be heuristically determined so that n_sim / n_draws is approximately equal to 20

Returns:

figplt.Figure

The figure object which can be readily saved to disk using fig.savefig().

Other Parameters:

conf_args

optional keyword arguments passed to the configurator

net_args

optional keyword arguments passed to the amortizer

plot_args

optional keyword arguments passed to plot_sbc()

load_pretrained_network()

Attempts to load a pre-trained network if checkpoint path is provided and a checkpoint manager exists.

train_online(epochs, iterations_per_epoch, batch_size, save_checkpoint=True, optimizer=None, reuse_optimizer=False, early_stopping=False, use_autograph=True, validation_sims=None, **kwargs)

Trains an amortizer via online learning. Additional keyword arguments are passed to the generative mode, configurator, and amortizer.

Parameters:

epochsint

Number of epochs (and number of times a checkpoint is stored)

iterations_per_epochint

Number of batch simulations to perform per epoch

batch_sizeint

Number of simulations to perform at each backprop step

save_checkpointbool, default: True

A flag to decide whether to save checkpoints after each epoch, if a checkpoint_path provided during initialization, otherwise ignored.

optimizertf.keras.optimizer.Optimizer or None

Optimizer for the neural network. None will result in tf.keras.optimizers.Adam using a learning rate of 5e-4 and a cosine decay from 5e-4 to 0. A custom optimizer will override default learning rate and schedule settings.

reuse_optimizerbool, optional, default: False

A flag indicating whether the optimizer instance should be treated as persistent or not. If False, the optimizer and its states are not stored after training has finished. Otherwise, the optimizer will be stored as ``self.optimizer` and re-used in further training runs.

early_stoppingbool, optional, default: False

Whether to use optional stopping or not during training. Could speed up training. Only works if validation_sims is not None, i.e., validation data has been provided.

use_autographbool, optional, default: True

Whether to use autograph for the backprop step. Could lead to enormous speed-ups but could also be harder to debug.

validation_simsdict or None, optional, default: None

Simulations used as a “validation set”. If dict, will assume it’s the output of a generative model and try amortizer.compute_loss(configurator(validation_sims)) after each epoch. If int, will assume it’s the number of sims to generate from the generative model before starting training. Only considered if a generative model has been provided during initialization. If None (default), no validation set will be used.

Returns:

lossesdict or pandas.DataFrame

A dictionary storing the losses across epochs and iterations

Other Parameters:

model_args

optional kwargs passed to the generative model

val_model_args:

optional kwargs passed to the generative model for generating validation data. Only useful if type(validation_sims) is int.

conf_args

optional kwargs passed to the configurator before each backprop (update) step.

val_conf_args

optional kwargs passed to the configurator then configuring the validation data.

net_args

optional kwargs passed to the amortizer

early_stopping_args

optional kwargs passed to the EarlyStopper

train_offline(simulations_dict, epochs, batch_size, save_checkpoint=True, optimizer=None, reuse_optimizer=False, early_stopping=False, validation_sims=None, use_autograph=True, **kwargs)

Trains an amortizer via offline learning. Assume parameters, data and optional context have already been simulated (i.e., forward inference has been performed).

Parameters:

simulations_dictdict

A dictionary containing the simulated data / context, if using the default keys, the method expects at least the mandatory keys sim_data and prior_draws to be present

epochsint

Number of epochs (and number of times a checkpoint is stored)

batch_sizeint

Number of simulations to perform at each backpropagation step

save_checkpointbool, default: True

Determines whether to save checkpoints after each epoch, if a checkpoint_path provided during initialization, otherwise ignored.

optimizertf.keras.optimizer.Optimizer or None

Optimizer for the neural network. None will result in tf.keras.optimizers.Adam using a learning rate of 5e-4 and a cosine decay from 5e-4 to 0. A custom optimizer will override default learning rate and schedule settings.

reuse_optimizerbool, optional, default: False

A flag indicating whether the optimizer instance should be treated as persistent or not. If False, the optimizer and its states are not stored after training has finished. Otherwise, the optimizer will be stored as self.optimizer and re-used in further training runs.

early_stoppingbool, optional, default: False

Whether to use optional stopping or not during training. Could speed up training. Only works if validation_sims is not None, i.e., validation data has been provided.

use_autographbool, optional, default: True

Whether to use autograph for the backprop step. Could lead to enormous speed-ups but could also be harder to debug.

validation_simsdict, int, or None, optional, default: None

Simulations used as a “validation set”. If dict, will assume it’s the output of a generative model and try amortizer.compute_loss(configurator(validation_sims)) after each epoch. If int, will assume it’s the number of sims to generate from the generative model before starting training. Only considered if a generative model has been provided during initialization. If None (default), no validation set will be used.

Returns:

lossesdict or pandas.DataFrame

A dictionary or a data frame storing the losses across epochs and iterations

Other Parameters:

val_model_args

optional kwargs passed to the generative model for generating validation data. Only useful if type(validation_sims) is int.

conf_args

optional kwargs passed to the configurator before each backprop (update) step.

val_conf_args

optional kwargs passed to the configurator then configuring the validation data.

net_args

optional kwargs passed to the amortizer

early_stopping_args

optional kwargs passed to the EarlyStopper

train_from_presimulation(presimulation_path, optimizer, save_checkpoint=True, max_epochs=None, reuse_optimizer=False, custom_loader=None, early_stopping=False, validation_sims=None, use_autograph=True, **kwargs)

Trains an amortizer via a modified form of offline training.

Like regular offline training, it assumes that parameters, data and optional context have already been simulated (i.e., forward inference has been performed).

Also like regular offline training, it is faster than online training in scenarios where simulations are slow. Unlike regular offline training, it uses each batch from the presimulated dataset only once during training, if not otherwise specified by a higher maximal number of epochs. Then, presimulated data is reused in a cyclic manner to achieve the desired number of epochs. A larger presimulated dataset is therefore required than for offline training, and the increase in speed gained by loading simulations instead of generating them on the fly comes at a cost: a large presimulated dataset takes up a large amount of hard drive space.

Parameters:

presimulation_pathstr

File path to the folder containing the files from the precomputed simulation. Ideally generated using a GenerativeModel’s presimulate_and_save method, otherwise must match the structure produced by that method. Each file contains the data for one epoch (i.e. a number of batches), and must be compatible with the custom_loader provided. The custom_loader must read each file into a collection (either a dictionary or a list) of simulation_dict objects. This is easily achieved with the pickle library: if the files were generated from collections of simulation_dict objects using pickle.dump, the _default_loader (default for custom_load) will load them using pickle.load. Training parameters like number of iterations and batch size are inferred from the files during training.

optimizertf.keras.optimizer.Optimizer

Optimizer for the neural network training. Since for this training, it is impossible to guess the number of iterations beforehead, an optimizer must be provided.

save_checkpointbool, optional, defaultTrue

Determines whether to save checkpoints after each epoch, if a checkpoint_path provided during initialization, otherwise ignored.

max_epochsint or None, optional, default: None

An optional parameter to limit or extend the number of epochs. If number of epochs is larger than the files of the dataset, presimulations will be reused.

reuse_optimizerbool, optional, default: False

A flag indicating whether the optimizer instance should be treated as persistent or not. If False, the optimizer and its states are not stored after training has finished. Otherwise, the optimizer will be stored as self.optimizer and re-used in further training runs.

custom_loadercallable, optional, default: self._default_loader

Must take a string file_path as an input and output a collection (dictionary or list) of simulation_dict objects. A simulation_dict has the keys prior_non_batchable_context, prior_batchable_context, prior_draws, sim_non_batchable_context, sim_batchable_context, and sim_data. Here, prior_draws and sim_data must have actual data as values, the rest are optional.

early_stoppingbool, optional, default: False

Whether to use optional stopping or not during training. Could speed up training.

validation_simsdict, int, or None, optional, default: None

Simulations used as a validation set. If dict, will assume it’s the output of a generative model and try amortizer.compute_loss(configurator(validation_sims)) after each epoch. If int, will assume it’s the number of sims to generate from the generative model before starting training. Only considered if a generative model has been provided during initialization. If None (default), no validation set will be used.

use_autographbool, optional, default: True

Whether to use autograph for the backprop step. Could lead to enormous speed-ups but could also be harder to debug.

Returns:

lossesdict or pandas.DataFrame

A dictionary or a data frame storing the losses across epochs and iterations

Other Parameters:

conf_args

optional keyword arguments passed to the configurator

net_args

optional keyword arguments passed to the amortizer

train_experience_replay(epochs, iterations_per_epoch, batch_size, save_checkpoint=True, optimizer=None, reuse_optimizer=False, buffer_capacity=1000, early_stopping=False, use_autograph=True, validation_sims=None, **kwargs)

Trains the network(s) via experience replay using a memory replay buffer, as utilized in reinforcement learning. Additional keyword arguments are passed to the generative mode, configurator, and amortizer. Read below for signature.

Parameters:

epochsint

Number of epochs (and number of times a checkpoint is stored)

iterations_per_epochint

Number of batch simulations to perform per epoch

batch_sizeint

Number of simulations to perform at each backpropagation step.

save_checkpointbool, optional, default: True

A flag to decide whether to save checkpoints after each epoch, if a checkpoint_path provided during initialization, otherwise ignored.

optimizertf.keras.optimizer.Optimizer or None

Optimizer for the neural network. None will result in tf.keras.optimizers.Adam using a learning rate of 5e-4 and a cosine decay from 5e-4 to 0. A custom optimizer will override default learning rate and schedule settings.

reuse_optimizerbool, optional, default: False

A flag indicating whether the optimizer instance should be treated as persistent or not. If False, the optimizer and its states are not stored after training has finished. Otherwise, the optimizer will be stored as self.optimizer and re-used in further training runs.

buffer_capacityint, optional, default: 1000

Max number of batches to store in buffer. For instance, if batch_size=32 and capacity_in_batches=1000, then the buffer will hold a maximum of 32 * 1000 = 32000 simulations. Be careful with memory! Important! Argument will be ignored if buffer has previously been initialized!

early_stoppingbool, optional, default: True

Whether to use optional stopping or not during training. Could speed up training. Only works if validation_sims is not None, i.e., validation data has been provided.

use_autographbool, optional, default: True

Whether to use autograph for the backprop step. Could lead to enormous speed-ups but could also be harder to debug.

validation_simsdict or None, optional, default: None

Simulations used as a “validation set”. If dict, will assume it’s the output of a generative model and try amortizer.compute_loss(configurator(validation_sims)) after each epoch. If int, will assume it’s the number of sims to generate from the generative model before starting training. Only considered if a generative model has been provided during initialization. If None (default), no validation set will be used.

Returns:

lossesdict or pandas.DataFrame

A dictionary or a data frame storing the losses across epochs and iterations.

Other Parameters:

model_args

optional kwargs passed to the generative model

val_model_args

optional kwargs passed to the generative model for generating validation data. Only useful if type(validation_sims) is int.

conf_args

optional kwargs passed to the configurator before each backprop (update) step.

val_conf_args

optional kwargs passed to the configurator then configuring the validation data.

net_args

optional kwargs passed to the amortizer

early_stopping_args:

optional kwargs passed to the EarlyStopper

train_rounds(rounds, sim_per_round, epochs, batch_size, save_checkpoint=True, optimizer=None, reuse_optimizer=False, early_stopping=False, use_autograph=True, validation_sims=None, **kwargs)

Trains an amortizer via round-based learning. In each round, sim_per_round data sets are simulated from the generative model and added to the data sets simulated in previous round. Then, the networks are trained for epochs on the augmented set of data sets.

Note

Training time will increase from round to round, since the number of simulations increases correspondingly. The final round will then train the networks on rounds * sim_per_round data sets, so make sure this number does not eat up all available memory.

Parameters:

roundsint

Number of rounds to perform (outer loop)

sim_per_roundint

Number of simulations per round

epochsint

Number of epochs (and number of times a checkpoint is stored, inner loop) within a round.

batch_sizeint

Number of simulations to use at each backpropagation step

save_checkpointbool, optional, default: True

A flag to decide whether to save checkpoints after each epoch, if a checkpoint_path provided during initialization, otherwise ignored.

optimizertf.keras.optimizer.Optimizer or None

Optimizer for the neural network training. None will result in tf.keras.optimizers.Adam using a learning rate of 5e-4 and a cosine decay from 5e-4 to 0. A custom optimizer will override default learning rate and schedule settings.

reuse_optimizerbool, optional, default: False

A flag indicating whether the optimizer instance should be treated as persistent or not. If False, the optimizer and its states are not stored after training has finished. Otherwise, the optimizer will be stored as self.optimizer and re-used in further training runs.

early_stoppingbool, optional, default: False

Whether to use optional stopping or not during training. Could speed up training. Only works if validation_sims is not None, i.e., validation data has been provided. Will be performed within rounds, not between rounds!

use_autographbool, optional, default: True

Whether to use autograph for the backprop step. Could lead to enormous speed-ups but could also be harder to debug.

validation_simsdict or None, optional, default: None

Simulations used as a “validation set”. If dict, will assume it’s the output of a generative model and try amortizer.compute_loss(configurator(validation_sims)) after each epoch. If int, will assume it’s the number of sims to generate from the generative model before starting training. Only considered if a generative model has been provided during initialization. If None (default), no validation set will be used.

Returns:

lossesdict or pandas.DataFrame

A dictionary or a data frame storing the losses across epochs and iterations

Other Parameters:

model_args

optional kwargs passed to the generative model

val_model_args

optional kwargs passed to the generative model for generating validation data. Only useful if type(validation_sims) is int.

conf_args

optional kwargs passed to the configurator before each backprop (update) step.

val_conf_args

optional kwargs passed to the configurator then configuring the validation data.

net_args

optional kwargs passed to the amortizer

early_stopping_args

optional kwargs passed to the EarlyStopper

mmd_hypothesis_test(observed_data, reference_data=None, num_reference_simulations=1000, num_null_samples=100, bootstrap=False)

Performs a sampling-based hypothesis test for detecting Out-Of-Simulation (model misspecification).

Parameters:

observed_datanp.ndarray

Observed data, shape (num_observed, …)

reference_datanp.ndarray

Reference data representing samples from the well-specified model, shape (num_reference, …)

num_reference_simulationsint, default: 1000

Number of reference simulations (M) simulated from the trainer’s generative model if no reference_data are provided.

num_null_samplesint, default: 100

Number of draws from the MMD sampling distribution under the null hypothesis “the trainer’s generative model is well-specified”

bootstrapbool, default: False

If true, the reference data (see above) are bootstrapped for each sample from the MMD sampling distribution. If false, a new data set is simulated for computing each draw from the MMD sampling distribution.

Returns:

mmd_null_samplesnp.ndarray

samples from the H0 sampling distribution (“well-specified model”)

mmd_observedfloat

summary MMD estimate for the observed data sets