Get Started¶
Installation¶
To install GMMVI (optionally) create a virtual environment and run
(.venv) $ pip install .
Usage¶
For performing the optimization, you can directly instantiate a GMMVI
and run GMMVI.train_iter()
in a loop, or, for adding basic logging
capability and easier integration, for example with WandB, you can instantiate a
GmmviRunner
and run
GmmviRunner.iterate_and_log(n)
in a loop.
Directly Using GMMVI¶
Before instantiating the GMMVI
, we need to create several other
objects, namely:
A
wrapped model
which stores the parameters of the GMM, as well as component-specific meta-information (reward histories, learning-rates, etc.)A
SampleDB
for storing samples.One object for each of the seven Design Choices.
Fortunately, each of these classes and also GMMVI
itself, have a
static method called build_from_config(), which allows to create
the object from a common config dictionary (which can be created from a YAML file). Using a common dictionary is
recommended, to ensure that the parameters passed to the different constructors are consistent (e.g. the sample
dimensions needs to be the same).
It is easiest to directly use GMMVI.build_from_config
,
which will automatically construct most of the required objects. However, you still need to pass
the dictionary containing the hyperparameters,
the
target distribution
,and the
initial model
.
The following example script directly uses GMMVI
using the hyperparameters
from the following YAML file: examples/example_config.yml.
import os
import logging
# Tensorflow may give warnings when the Cholesky decomposition fails.
# However, these warning can usually be ignored because the NgBasedOptimizer
# will handle them by rejecting the update and decreasing the stepsize for
# the failing component. To keep the console uncluttered, we suppress warnings.
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # ERROR
logging.getLogger('tensorflow').setLevel(logging.ERROR)
import tensorflow as tf
from gmmvi.optimization.gmmvi import GMMVI
from gmmvi.configs import load_yaml
from gmmvi.experiments.target_distributions.logistic_regression import make_breast_cancer
from gmmvi.models.full_cov_gmm import FullCovGMM
from gmmvi.models.gmm_wrapper import GmmWrapper
from gmmvi.experiments.setup_experiment import construct_initial_mixture
#For creating a GMMVI object using GMMVI.build_from_config, we need:
# 1. A dictionary containing the hyperparameters
my_path = os.path.dirname(os.path.realpath(__file__))
config = load_yaml(os.path.join(my_path, "example_config.yml"))
# 2. A target distribution
target_distribution = make_breast_cancer()
# 3. An (wrapped) initial model
dims = target_distribution.get_num_dimensions()
initial_weights = tf.ones(1, tf.float32)
initial_means = tf.zeros((1, dims), tf.float32)
initial_covs = tf.reshape(100 * tf.eye(dims), [1, dims, dims])
model = FullCovGMM(initial_weights, initial_means, initial_covs)
# Above config contains a section model_initialization, and, therefore,
# we could also create the initial model using:
# model = construct_initial_mixture(dims, **config["model_initialization"])
wrapped_model = GmmWrapper.build_from_config(model=model, config=config)
# Now we can create the GMMVI object and start optimizing
gmmvi = GMMVI.build_from_config(config=config,
target_distribution=target_distribution,
model=wrapped_model)
max_iter = 1001
for n in range(max_iter):
gmmvi.train_iter()
if n % 100 == 0:
samples = gmmvi.model.sample(1000)[0]
elbo = tf.reduce_mean(target_distribution.log_density(samples)
- model.log_density(samples))
print(f"{n}/{max_iter}: "
f"The model now has {gmmvi.model.num_components} components "
f"and an elbo of {elbo}.")
The script can be found under examples/1_directly_using_gmmvi.py.
Using the GmmviRunner¶
The GmmviRunner
wraps around
GMMVI
to add logging capabilities.
Furthermore, the GmmviRunner
takes care
of initializing the model and the target distribution (when using one of the provided
target distributions). Hence, we only need to provide the config to create it,
as shown by the following script:
import os
from gmmvi.gmmvi_runner import GmmviRunner
from gmmvi.configs import load_yaml
my_path = os.path.dirname(os.path.realpath(__file__))
config = load_yaml(os.path.join(my_path, "example_config.yml"))
gmmvi_runner = GmmviRunner.build_from_config(config)
for n in range(10001):
gmmvi_runner.iterate_and_log(n)
The script can be found under examples/2_using_the_gmmvi_runner.py.
Using the GmmviRunner with Default Configs¶
We can also directly create a default config based on the 7-letter Codeword to specify the design choices, thereby, not requiring an external YAML file:
from gmmvi.gmmvi_runner import GmmviRunner
import gmmvi.configs
# In this example, we will create the config for a GmmviRunner using default configs
# for a given Codename (we weill use SAMYROX) and an and an environment name
# (we will use GMM20).
# Let's first get the default config for SAMYROX
algorithm_config = gmmvi.configs.get_default_algorithm_config("SAMYROX")
# Internally, this loaded the yaml files in gmmvi/configs/module_configs corresponding
# to the chosen design choices and stored them in a single dict "algorithm_config".
# Note that these default values were chosen independently for every design choice,
# and, thus, may not always be sensible. For example, the initial_stepsize defined in
# gmmvi/configs/module_configs/component_stepsize_adaptation/improvement_based.yml
# (Codeletter "R") is suitable if the stepsize is treated as a trust-region
# (Codeletter "T"), but not if it directly corresponds to the stepsize
# (Codeletter "I" or "Y")! Hence, we will overwrite the stepsize to something more
# suitable for SAMYROX:
better_stepsize_config = {
'initial_stepsize': 0.0001,
'min_stepsize': 0.0001,
'max_stepsize': 0.001
}
algorithm_config = gmmvi.configs.update_config(algorithm_config, better_stepsize_config)
# We will use a target distribution that was shipped with the framework, namely "gmm20":
environment_config = gmmvi.configs.get_default_experiment_config("gmm20")
# The last call searched configs/experiment_configs for a corresponding yml-file and found
# gmm20.yml and stored the config in the dictionary "environment_config". We now just need
# to merge both config files:
config = gmmvi.configs.update_config(algorithm_config, environment_config)
# Create the GmmviRunner and start optimizing.
gmmvi_runner = GmmviRunner.build_from_config(config=config)
for n in range(1500):
gmmvi_runner.iterate_and_log(n)
The script can be found under examples/3_gmmvi_runner_with_default_configs.py.
Using the GmmviRunner with Custom Environments¶
We can still use the GmmviRunner
with custom environments, but we need to store the
target distribution object
in the config:
from gmmvi.gmmvi_runner import GmmviRunner
from gmmvi.configs import get_default_algorithm_config, update_config
import tensorflow as tf
import numpy as np
import matplotlib
matplotlib.use("tkAgg")
import matplotlib.pyplot as plt
# For creating a custom environment, we need to extend
# gmmvi.experiments.target_distributions.lnpdf.LNPDF:
from gmmvi.experiments.target_distributions.lnpdf import LNPDF
class Rosenbrock(LNPDF):
""" We treat the negative Rosenbrock function as unnormalized target distribution.
We implement it in numpy and do not allow GMMVI to backpropagate through log_density().
As we want to use Stein's Lemma for estimating the natural gradient (Codeletter "S"),
we need to implement the gradient ourselves, and, therefore, we set
use_log_density_and_grad=True and implement the corresponding method.
"""
def __init__(self):
super(Rosenbrock, self).__init__(use_log_density_and_grad=True,
safe_for_tf_graph=False)
self.a = 1
self.b = 100
def get_num_dimensions(self) -> int:
return 2
def log_density(self, samples: tf.Tensor) -> tf.Tensor:
x = samples[:, 0].numpy().astype(np.float32)
y = samples[:, 1].numpy().astype(np.float32)
my_log_density = -((self.a - x)**2 + self.b * (y - x**2)**2)
return tf.convert_to_tensor(my_log_density, dtype=tf.float32)
def log_density_and_grad(self, samples: tf.Tensor) -> tf.Tensor:
x = samples[:, 0].numpy().astype(np.float32)
y = samples[:, 1].numpy().astype(np.float32)
my_log_density = -((self.a - x)**2 + self.b * (y - x**2)**2)
my_grad_x = -(-2 * (self.a - x) - 4 * self.b * (y - x**2) * x)
my_grad_y = -(2 * self.b * (y - x**2))
my_grad = np.vstack((my_grad_x, my_grad_y)).T
return [tf.convert_to_tensor(my_log_density, dtype=tf.float32),
tf.convert_to_tensor(my_grad, dtype=tf.float32)]
# We can also use the GmmviRunner, when using custom environments, but we have
# to put the LNPDF object into the dict. Furthermore, we need to define the other
# environment-specific settings that would otherwise be defined in
# the corresponding config in gmmvi/config/experiment_configs:
environment_config = {
"target_fn": Rosenbrock(),
"start_seed": 0,
"environment_name": "Rosenbrock",
"model_initialization": {
"use_diagonal_covs": False,
"num_initial_components": 1,
"prior_mean": 0.,
"prior_scale": 1.,
"initial_cov": 1.,
},
"gmmvi_runner_config": {
"log_metrics_interval": 100
},
"use_sample_database": True,
"max_database_size": int(1e6),
"temperature": 1.
}
# We will again use the automatically generated config for the algorithm,
# but this time, we will use "SAMTRUX". The default settings are reasonable for
# SAMTRUX, so we do not make any modifications to the hyperparameters.
algorithm_config = get_default_algorithm_config("SAMTRUX")
# Now we just need to merge the configs and use GmmviRunner as before:
merged_config = update_config(algorithm_config, environment_config)
gmmvi_runner = GmmviRunner.build_from_config(merged_config)
for n in range(500):
gmmvi_runner.iterate_and_log(n)
# Plot samples from our "Rosenbrock-distribution"
test_samples = gmmvi_runner.gmmvi.model.sample(10000)[0]
plt.plot(test_samples[:, 0], test_samples[:, 1], 'x')
plt.show()
plt.pause(0.1)
The script can be found under examples/4_gmmvi_runner_with_custom_environments.py.