Source code for gmmvi.optimization.gmmvi_modules.weight_stepsize_adaptation

import tensorflow as tf

from gmmvi.models.gmm_wrapper import GmmWrapper


[docs]class WeightStepsizeAdaptation: """ This class provides a common interface for adapting the stepsize for the weight update. There are currently three options for weight stepsize adpatation: 1. The :py:class:`FixedWeightStepsizeAdaptation` is a dummy-class, that does not do anything. 2. The :py:class:`DecayingWeightStepsizeAdaptation` uses exponential decay. 3. The :py:class:`ImprovementBasedWeightStepsizeAdaptation` uses a procedure similar to VIPS :cite:p:`Arenz2020` to increase the stepsize if the mixture improved during the last updates, and to decrease it otherwise. Parameters: initial_stepsize: float The initial stepsize for the weight update. """ def __init__(self, initial_stepsize: tf.float32): self.stepsize = tf.Variable(initial_stepsize, dtype=tf.float32)
[docs] @staticmethod def build_from_config(config, gmm_wrapper): """This static method provides a convenient way to create a :py:class:`FixedWeightStepsizeAdaptation`, :py:class:`DecayingWeightStepsizeAdaptation` or :py:class:`ImprovementBasedWeightStepsizeAdaptation` 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. """ if config["weight_stepsize_adapter_type"] == "fixed": return FixedWeightStepsizeAdaptation(**config['weight_stepsize_adapter_config']) elif config["weight_stepsize_adapter_type"] == "decaying": return DecayingWeightStepsizeAdaptation(**config['weight_stepsize_adapter_config']) elif config["weight_stepsize_adapter_type"] == "improvement_based": return ImprovementBasedWeightStepsizeAdaptation(gmm_wrapper, **config['weight_stepsize_adapter_config']) else: raise ValueError( f"config['weight_stepsize_adapter_type'] is '{config['weight_stepsize_adapter_type']}' " f"which is an unknown type")
def _update_stepsize(self): pass
[docs] def update_stepsize(self): """ Update the stepsizes, according to the chosen procedure. Returns: float: the updated stepsize. """ self._update_stepsize() return self.stepsize
[docs]class FixedWeightStepsizeAdaptation(WeightStepsizeAdaptation): """ This class is a dummy class, that can be used when we want to keep the stepsize for the weight update constant. Parameters: initial_stepsize: float The initial stepsize for the weight update. """ def __init__(self, initial_stepsize: tf.float32): super(FixedWeightStepsizeAdaptation, self).__init__(initial_stepsize)
[docs]class DecayingWeightStepsizeAdaptation(WeightStepsizeAdaptation): """ This class implements an exponentially decaying stepsize schedule. See :cite:p:`Khan2018a`. Parameters: gmm_wrapper: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The wrapped. annealing_exponent: float controls how fast the stepsize decays initial_stepsize: float The initial stepsize for the weight update. """ def __init__(self, initial_stepsize: tf.float32, annealing_exponent: tf.float32): super(DecayingWeightStepsizeAdaptation, self).__init__(initial_stepsize) self.initial_stepsize = tf.constant(initial_stepsize) self.annealing_exponent = annealing_exponent self.num_weight_updates = tf.Variable(0, dtype=tf.float32) def _update_stepsize(self): """ Updates the stepsize using exponential decay. More specifially, the new stepsize is given by :math:`\\frac{\\text{initial\\_stepsize}}{1 + \\text{num\\_iterations}^\\text{annealing\\_exponent}}`. Returns: float: the updated stepsize. """ self.stepsize.assign(self.initial_stepsize / (1. + tf.math.pow(self.num_weight_updates, self.annealing_exponent))) self.num_weight_updates.assign_add(1.)
[docs]class ImprovementBasedWeightStepsizeAdaptation(WeightStepsizeAdaptation): """ Increases the stepsize if the last weight update increased its reward, decreases it otherwise. Parameters: gmm_wrapper: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The wrapped model. initial_stepsize: float The initial stepsize for the weight update. min_stepsize: float Do not not decrease the stepsize below this point max_stepsize: float Do not increase the stepsize above this point stepsize_inc_factor: float Factor (>1) for increasing the stepsize stepsize_dec_factor: float Factor in the range [0, 1] for decreasing the stepsize """ def __init__(self, model: GmmWrapper, initial_stepsize: tf.float32, min_stepsize: tf.float32, max_stepsize: tf.float32, stepsize_inc_factor: tf.float32, stepsize_dec_factor: tf.float32): super(ImprovementBasedWeightStepsizeAdaptation, self).__init__(initial_stepsize) self.initial_stepsize = tf.constant(initial_stepsize) self.model = model self.min_stepsize = min_stepsize self.max_stepsize = max_stepsize self.stepsize_inc_factor = stepsize_inc_factor self.stepsize_dec_factor = stepsize_dec_factor self.elbo_history = tf.Variable([tf.float32.min], shape=[None], dtype=tf.float32) def _update_stepsize(self): """ Updates the stepsize of each component based on the mixture model improvement. Returns: float: the updated stepsize. """ elbo = tf.reduce_sum(self.model.weights * self.model.reward_history[:,-1]) - tf.reduce_sum(self.model.weights * self.model.log_weights) self.elbo_history.assign(tf.concat((self.elbo_history, tf.expand_dims(elbo,0)), axis=0)) if self.elbo_history[-1] > self.elbo_history[-2]: self.stepsize.assign(tf.math.minimum( self.stepsize_inc_factor * self.stepsize, self.max_stepsize)) else: self.stepsize.assign(tf.math.maximum( self.stepsize_dec_factor * self.stepsize, self.min_stepsize))