Source code for gmmvi.optimization.gmmvi_modules.sample_selector

import tensorflow as tf
from gmmvi.optimization.sample_db import SampleDB
from gmmvi.models.gmm_wrapper import GmmWrapper
from gmmvi.experiments.target_distributions.lnpdf import LNPDF

[docs]class SampleSelector: """ Provides the interface for selecting samples for performing the updates at the beginning of every iteration. The samples are evaluated on the target distribution and used for updating the weights, means and covariance of the GMM. There are currently two options for estimating the natural gradient: 1. The :py:class:`VipsSampleSelector` use the procedure described by :cite:t:`Arenz2018,Arenz2020` to ensure that we have samples in the vicinity of every component, enabling us to perform a stable update on every component. 2. The :py:class:`LinSampleSelector` uses the procedure described by :cite:t:`Lin2019` which draws samples according to the weights of the current mixture model, aiming for better sample efficiency. Parameters: target_distribution: :py:class:`LNPDF` The target distribution is used for evaluating the newly drawn samples. model: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The wrapped model is used for drawing the samples. sample_db: :py:class:`SampleDB` The new samples and their target_densities (and gradients) are stored in the sample database. """ def __init__(self, target_distribution: LNPDF, model: GmmWrapper, sample_db: SampleDB): self.target_distribution = target_distribution self.model = model self.sample_db = sample_db if self.target_distribution.safe_for_tf_graph: self.get_target_grads = tf.function(self.get_target_grads, experimental_relax_shapes=True)
[docs] @staticmethod def build_from_config(config, gmm_wrapper, sample_db, target_distribution): """This static method provides a convenient way to create a :py:class:`VipsSampleSelector`, or :py:class:`LinSampleSelector` depending on the provided config. Parameters: config: dict The dictionary is typically read from YAML a file, and holds all hyperparameters. gmm_wrapper: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The wrapped model is used for drawing the samples. sample_db: :py:class:`SampleDB` The new samples and their target_densities (and gradients) are stored in the sample database. target_distribution: :py:class:`LNPDF` The target distribution is used for evaluating the newly drawn samples. """ if config["sample_selector_type"] == "component-based": return VipsSampleSelector(target_distribution, gmm_wrapper, sample_db, **config['sample_selector_config']) elif config["sample_selector_type"] == "mixture-based": return LinSampleSelector(target_distribution, gmm_wrapper, sample_db, **config['sample_selector_config']) else: raise ValueError( f"config['sample_selector_type'] is '{config['sample_selector_type']}' which is an unknown type")
def target_uld(self, samples: tf.Tensor) -> tf.Tensor: return self.target_distribution.log_density(samples) def get_target_grads(self, samples: tf.Tensor) -> [tf.Tensor, tf.Tensor]: if self.target_distribution.use_log_density_and_grad: # useful if we can't backprop through target_uld target, gradient = self.target_distribution.log_density_and_grad(samples) else: with tf.GradientTape(persistent=False) as tape: tape.watch(samples) target = self.target_distribution.log_density(samples) gradient = tape.gradient(target, samples) return gradient, target
[docs] def select_samples(self) -> [tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]: """ Select the samples for current learning iteration and stores the data in the sample database. Returns: tuple(tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor): **samples** - a tensor of shape number_of_selected_samples x number_of_dimensions **old_samples_pdf** - a tensor of shape number_of_selected_samples, containing the log-densities of the distribution that was effectively used to obtain the selected samples. Needed for importance weighting. **target_lnpdfs** - a tensor of shape number_of_selected_samples, containing the log-densities of the target distrbution for each selected sample, :math:`\\log p(\\mathbf{x})`. **target_grads** - a tensor of shape number_of_selected_samples x num_dimensions, containing the gradients of the log-densities of the target distrbution for each selected sample, :math:`\\nabla_{\\mathbf{x}} \\log p(\\mathbf{x})`. """ raise NotImplementedError
[docs]class VipsSampleSelector(SampleSelector): """ Selects the samples according to the procedure described by :cite:t:`Arenz2018,Arenz2020`. This class uses the procedure described by :cite:t:`Arenz2018,Arenz2020` to ensure that we have samples in the vicinity of every component. It uses two passes. In the first pass, it selects a given number of samples from the sample database. In the second pass, it computes the effective sample size for every component (based on the importance weights) and compares the effective sample size with a given desired number of samples. It then draws from every component the respective missing number of samples. Parameters: target_distribution: :py:class:`LNPDF` The target distribution is used for evaluating the newly drawn samples. model: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The wrapped model is used for drawing the samples. sample_db: :py:class:`SampleDB` The database is used for reusing samples from previous iterations and for storing the new samples and their target_densities (and gradients). desired_samples_per_component: int The desired number of samples for every component. ratio_reused_samples_to_desired: float In the first pass, we reuse the number_of_components * *ratio_reused_samples_to_desired* * *desired_samples_per_component* freshest samples from the database. """ def __init__(self, target_distribution: LNPDF, model: GmmWrapper, sample_db: SampleDB, desired_samples_per_component: int, ratio_reused_samples_to_desired: float): super(VipsSampleSelector, self).__init__(target_distribution, model, sample_db) self.sample_db = sample_db self.desired_samples_per_component = desired_samples_per_component self.reused_samples_per_component = tf.cast(tf.math.floor( ratio_reused_samples_to_desired * desired_samples_per_component), dtype=tf.int32)
[docs] def get_effective_samples(self, model_densities: tf.Tensor, oldsamples_pdf: tf.Tensor) -> tf.Tensor: """ Computes the effective sample size based on the log-densities of the target distribution and the log-densities of the background distribution. Parameters: model_densities: tf.Tensor The log-densities of the individual components, :math:`\\log q(\\mathbf{x}|o)` oldsamples_pdf: tf.Tensor The log-densities of the distribution that was effectively used for obtaining the selected samples Returns: float: the effective number of samples """ log_weight = model_densities - tf.expand_dims(oldsamples_pdf, axis=0) log_weight = log_weight - tf.reduce_logsumexp(log_weight, axis=1, keepdims=True) weights = tf.exp(log_weight) num_effective_samples = 1. / tf.reduce_sum(weights * weights, axis=1) return num_effective_samples
[docs] def sample_where_needed(self, samples: tf.Tensor, oldsamples_pdf: tf.Tensor, num_desired_samples: int = None)\ -> [tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]: """ Computes the components' effective sample sizes for the given set of samples and draws, for every component i, :math:`n_{\\text{des}} - n_{\\text{eff,i}}` new samples. Parameters: samples: tf.Tensor the samples that were chosen during the first pass oldsamples_pdf: tf.Tensor The log-densities of the distribution that was effectively used for obtaining the selected samples num_desired_samples: int The number of desired samples per component Returns: tuple(tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor): **new_samples** - a tf.Tensor, the newly drawn samples **new_target_lnpdfs** - a tf.Tensor, the log-densities of the target distribution on the newly drawn samples, :math:`\\log p(\\mathbf{x})`. **new_target_grads** - a tf.Tensor, the gradients of the log-densities for the newly drawn samples, :math:`\\nabla_{\\mathbf{x}} \\log p(\\mathbf{x})`. **mapping** - a tf.Tensor, for every sample the one-dimensional tensor contains the index of the component that was used for drawing that sample. """ if num_desired_samples is None: num_desired_samples = self.desired_samples_per_component if tf.shape(samples)[0] == 0: num_effective_samples = tf.zeros((self.model.num_components), dtype=tf.int32) else: model_logpdfs = self.model.component_log_densities(samples) num_effective_samples = tf.cast(tf.math.floor(self.get_effective_samples(model_logpdfs, oldsamples_pdf)), dtype=tf.int32) num_additional_samples = tf.math.maximum(1, num_desired_samples - num_effective_samples) new_samples, mapping = self.model.sample_from_components_no_shuffle(num_additional_samples) new_target_grads, new_target_lnpdfs = self.get_target_grads(new_samples) return new_samples, new_target_lnpdfs, new_target_grads, mapping
[docs] def select_samples(self) -> [tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]: # Get old samples from the database num_samples_to_reuse = self.reused_samples_per_component * self.model.num_components oldsamples_pdf, samples, _, _, _ = self.sample_db.get_newest_samples(num_samples_to_reuse) num_reused_samples = tf.shape(samples)[0] # Get additional samples to ensure a desired effective sample size for every component new_samples, new_target_lnpdfs, new_target_grads, mapping = self.sample_where_needed(samples, oldsamples_pdf) self.sample_db.add_samples(new_samples, self.model.means, self.model.chol_cov, new_target_lnpdfs, new_target_grads, mapping) num_new_samples = tf.shape(new_samples)[0] # We call get_newest_samples again in order to recompute the background densities oldsamples_pdf, samples, mapping, target_lnpdfs, target_grads = self.sample_db.get_newest_samples(num_reused_samples + num_new_samples) return samples, mapping, oldsamples_pdf, target_lnpdfs, target_grads
[docs]class LinSampleSelector(SampleSelector): """ Selects the samples according to the procedure described by :cite:t:`Lin2019`. This class uses the procedure described by :cite:t:`Lin2019` by drawing new samples for the current mixture model. We also implemented the two-phase procedure of the :py:class:`VipsSampleSelector` to reuse samples from the database and redraw samples based on a desired number of samples. However, in contrast to the :py:class:`VipsSampleSelector`, we compute the effective sample size not per component, but for the whole mixture, and redraw samples n_eff - *desired_samples_per_component* new samples from the mixture model. The exact procedure of :cite:t:`Lin2019` can be reproduced, when choosing *ratio_reused_samples_to_desired* = 0, where always a fixed number of new samples is drawn from the mixture model. Parameters: target_distribution: :py:class:`LNPDF` The target distribution is used for evaluating the newly drawn samples. model: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The wrapped model is used for drawing the samples. sample_db: :py:class:`SampleDB` The database is used for reusing samples from previous iterations and for storing the new samples and their target_densities (and gradients). desired_samples_per_component: int The desired number for the mixture update. ratio_reused_samples_to_desired: float In the first pass, we reuse the *ratio_reused_samples_to_desired* * *desired_samples_per_component* freshest samples from the database. """ def __init__(self, target_distribution: LNPDF, model: GmmWrapper, sample_db: SampleDB, desired_samples_per_component: int, ratio_reused_samples_to_desired: float): super(LinSampleSelector, self).__init__(target_distribution, model, sample_db) self.desired_samples_per_component = desired_samples_per_component self.reused_samples_per_component = tf.cast(tf.math.floor( ratio_reused_samples_to_desired * desired_samples_per_component), dtype=tf.int32)
[docs] def get_effective_samples(self, model_densities: tf.Tensor, oldsamples_pdf: tf.Tensor) -> tf.Tensor: """ Computes the effective sample size of the mixture model based on the log-densities of the target distribution and the log-densities of the background distribution. Parameters: model_densities: tf.Tensor The log-densities of the mixture model, :math:`\\log q(\\mathbf{x})`. oldsamples_pdf: tf.Tensor The log-densities of the distribution that was effectively used for obtaining the selected samples Returns: float: the effective number of samples """ log_weight = model_densities - tf.expand_dims(oldsamples_pdf, axis=0) log_weight = log_weight - tf.reduce_logsumexp(log_weight, axis=1, keepdims=True) weights = tf.exp(log_weight) num_effective_samples = 1. / tf.reduce_sum(weights * weights, axis=1) return num_effective_samples
[docs] @tf.function def sample_where_needed(self) -> [tf.Tensor, tf.Tensor, int]: """ Computes the mixture model's effective sample size for the given set of samples and draws :math:`n_{\\text{des}} - n_{\\text{eff}}` new samples from the mixture model. Parameters: samples: tf.Tensor the samples that were chosen during the first pass oldsamples_pdf: tf.Tensor The log-densities of the distribution that was effectively used for obtaining the selected samples num_desired_samples: int The number of desired samples per component Returns: tuple(tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor): **new_samples** - a tensor containing the newly drawn samples **new_target_lnpdfs** - a tensor containing the log-densities of the target distribution on the newly drawn samples, :math:`\\log p(\\mathbf{x})`. **new_target_grads** - a tensor containing the gradients of the log-densities for the newly drawn samples. **mapping** - a tensor containing for every sample the one-dimensional tensor contains the index of the component that was used for drawing that sample. """ # Get old samples from the database num_samples_to_reuse = self.reused_samples_per_component * self.model.num_components oldsamples_pdf, old_samples, _, _, _ = self.sample_db.get_newest_samples(num_samples_to_reuse) num_reused_samples = tf.shape(old_samples)[0] # Get additional samples to ensure a desired effective sample size for every component if tf.shape(old_samples)[0] == 0: num_effective_samples = tf.zeros((1), dtype=tf.int32) else: model_logpdfs = self.model.log_density(old_samples) num_effective_samples = tf.cast(tf.math.floor(self.get_effective_samples(model_logpdfs, oldsamples_pdf)), dtype=tf.int32) num_additional_samples = tf.math.maximum(1, self.desired_samples_per_component - num_effective_samples) new_samples, mapping = self.model.sample(tf.squeeze(num_additional_samples)) return new_samples, mapping, num_reused_samples
[docs] def select_samples(self) -> [tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]: # Get additional samples to ensure a desired effective sample size for every component new_samples, mapping, num_reused_samples = self.sample_where_needed() new_target_grads, new_target_lnpdfs = self.get_target_grads(new_samples) self.sample_db.add_samples(new_samples, self.model.means, self.model.chol_cov, new_target_lnpdfs, new_target_grads, mapping) # We call get_newest_samples again in order to recompute the background densities samples_this_iter = num_reused_samples + tf.shape(new_samples)[0] oldsamples_pdf, samples, mapping, target_lnpdfs, target_grads = self.sample_db.get_newest_samples(samples_this_iter) return samples, mapping, oldsamples_pdf, target_lnpdfs, target_grads