Module WUT.Gaussian_Mixtures
Expand source code
import numpy as np
import tensorflow as tf
import types
from tensorflow.python.keras.layers.ops import core as core_ops
import mdn
from tensorflow.compat.v1.keras import layers
from tensorflow.python.keras import activations
from tensorflow_probability import distributions as tfd
from keras import backend as K
from keras import activations, initializers
import tensorflow_probability as tfp
class Gaussian_Mixtures():
def __init__(self, Model, num_mixtures=1):
"""Gaussian_MIxtures Initializer. Turns a neural network into an GMN.
Args:
Model: Input Keras Model.
num_mixtures: how many total gaussians would you like to fit the output space to.
Returns:
Nothing lol
"""
self.model = Model()
layer = self.model.layers[-1]
self.output_dim = layer.units
layer.output_dim = layer.units
self.num_mix = num_mixtures
layer.num_mix = num_mixtures
with tf.name_scope('MDN'):
layer.mdn_mus = layers.Dense(layer.num_mix * layer.output_dim, name='mdn_mus') # mix*output vals, no activation
layer.mdn_sigmas = layers.Dense(self.num_mix * self.output_dim, activation=self.elu_plus_one_plus_epsilon, name='mdn_sigmas') # mix*output vals exp activation
layer.mdn_pi = layers.Dense(self.num_mix, name='mdn_pi') # mix vals, logits
def build(self, input_shape):
with tf.name_scope('mus'):
self.mdn_mus.build(input_shape)
with tf.name_scope('sigmas'):
self.mdn_sigmas.build(input_shape)
with tf.name_scope('pis'):
self.mdn_pi.build(input_shape)
def call_func(self, x):
with tf.name_scope('MDN'):
mdn_out = layers.concatenate([self.mdn_mus(x),
self.mdn_sigmas(x),
self.mdn_pi(x)],
name='mdn_outputs')
return mdn_out
def compute_output_shape(self, input_shape):
"""Returns output shape, showing the number of mixture parameters."""
return (input_shape[0], (2 * self.output_dim * self.num_mix) + self.num_mix)
def get_config(self):
config = {
"output_dimension": self.output_dim,
"num_mixtures": self.num_mix
}
base_config = super(Dense, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
layer.build = types.MethodType(build, layer)
layer.call = types.MethodType(call_func, layer)
layer._trainable_weights = layer.mdn_mus.trainable_weights + layer.mdn_sigmas.trainable_weights + layer.mdn_pi.trainable_weights
layer._non_trainable_weights = layer.mdn_mus.non_trainable_weights + layer.mdn_sigmas.non_trainable_weights + layer.mdn_pi.non_trainable_weights
layer.compute_output_shape = types.MethodType(compute_output_shape, layer)
layer.get_config = types.MethodType(get_config, layer)
def elu_plus_one_plus_epsilon(self, x):
"""ELU activation with a very small addition to help prevent
NaN in loss."""
return tf.keras.backend.elu(x) + 1 + .00001
def compile(self, *args, loss=None, **kwargs):
"""compile. Literally use this as you'd use the normal compile, but don't use your own loss functions, unless your really deep in this. Let the default one go
Args:
loss: if you really wanna make your own loss function
Returns:
Nothing lol
"""
if loss is None:
loss = mdn.get_mixture_loss_func(self.output_dim, self.num_mix)
kwargs['loss'] = loss
self.model.compile(*args, **kwargs)
def fit(self, *args, **kwargs):
"""fit. Literally use this as you'd use the normal keras fit.
Returns:
Nothing lol.
"""
self.model.fit(*args, **kwargs)
def evaluate(self, *args, **kwargs):
"""evaluate. Literally use this as you'd use the normal keras evaluate.
Args:
Returns:
The model's score, evaluated on whatever inputs you just fed it.
"""
return self.model.evaluate( *args, **kwargs)
def predict(self, *args, **kwargs):
"""predict. Literally use this as you'd use the normal keras predict.
Returns:
a probability distribution over outputs, for each input.
"""
all_preds = self.model.predict(*args, **kwargs)
return self.get_dist(all_preds)
def get_dist(self, y_pred):
"""turns an output into a distribution. Literally use this as you'd use the normal keras predict.
Args:
y_pred: nn output
Returns:
a probability distribution over outputs, for each input.
"""
num_mix = self.num_mix
output_dim = self.output_dim
y_pred = tf.reshape(y_pred, [-1, (2 * num_mix * output_dim) + num_mix], name='reshape_ypreds')
out_mu, out_sigma, out_pi = tf.split(y_pred, num_or_size_splits=[num_mix * output_dim,
num_mix * output_dim,
num_mix],
axis=1, name='mdn_coef_split')
cat = tfd.Categorical(logits=out_pi)
component_splits = [output_dim] * num_mix
mus = tf.split(out_mu, num_or_size_splits=component_splits, axis=1)
sigs = tf.split(out_sigma, num_or_size_splits=component_splits, axis=1)
coll = [tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale) for loc, scale
in zip(mus, sigs)]
return tfd.Mixture(cat=cat, components=coll)
Classes
class Gaussian_Mixtures (Model, num_mixtures=1)-
Gaussian_MIxtures Initializer. Turns a neural network into an GMN.
Args
Model- Input Keras Model.
num_mixtures- how many total gaussians would you like to fit the output space to.
Returns
Nothing lol
Expand source code
class Gaussian_Mixtures(): def __init__(self, Model, num_mixtures=1): """Gaussian_MIxtures Initializer. Turns a neural network into an GMN. Args: Model: Input Keras Model. num_mixtures: how many total gaussians would you like to fit the output space to. Returns: Nothing lol """ self.model = Model() layer = self.model.layers[-1] self.output_dim = layer.units layer.output_dim = layer.units self.num_mix = num_mixtures layer.num_mix = num_mixtures with tf.name_scope('MDN'): layer.mdn_mus = layers.Dense(layer.num_mix * layer.output_dim, name='mdn_mus') # mix*output vals, no activation layer.mdn_sigmas = layers.Dense(self.num_mix * self.output_dim, activation=self.elu_plus_one_plus_epsilon, name='mdn_sigmas') # mix*output vals exp activation layer.mdn_pi = layers.Dense(self.num_mix, name='mdn_pi') # mix vals, logits def build(self, input_shape): with tf.name_scope('mus'): self.mdn_mus.build(input_shape) with tf.name_scope('sigmas'): self.mdn_sigmas.build(input_shape) with tf.name_scope('pis'): self.mdn_pi.build(input_shape) def call_func(self, x): with tf.name_scope('MDN'): mdn_out = layers.concatenate([self.mdn_mus(x), self.mdn_sigmas(x), self.mdn_pi(x)], name='mdn_outputs') return mdn_out def compute_output_shape(self, input_shape): """Returns output shape, showing the number of mixture parameters.""" return (input_shape[0], (2 * self.output_dim * self.num_mix) + self.num_mix) def get_config(self): config = { "output_dimension": self.output_dim, "num_mixtures": self.num_mix } base_config = super(Dense, self).get_config() return dict(list(base_config.items()) + list(config.items())) layer.build = types.MethodType(build, layer) layer.call = types.MethodType(call_func, layer) layer._trainable_weights = layer.mdn_mus.trainable_weights + layer.mdn_sigmas.trainable_weights + layer.mdn_pi.trainable_weights layer._non_trainable_weights = layer.mdn_mus.non_trainable_weights + layer.mdn_sigmas.non_trainable_weights + layer.mdn_pi.non_trainable_weights layer.compute_output_shape = types.MethodType(compute_output_shape, layer) layer.get_config = types.MethodType(get_config, layer) def elu_plus_one_plus_epsilon(self, x): """ELU activation with a very small addition to help prevent NaN in loss.""" return tf.keras.backend.elu(x) + 1 + .00001 def compile(self, *args, loss=None, **kwargs): """compile. Literally use this as you'd use the normal compile, but don't use your own loss functions, unless your really deep in this. Let the default one go Args: loss: if you really wanna make your own loss function Returns: Nothing lol """ if loss is None: loss = mdn.get_mixture_loss_func(self.output_dim, self.num_mix) kwargs['loss'] = loss self.model.compile(*args, **kwargs) def fit(self, *args, **kwargs): """fit. Literally use this as you'd use the normal keras fit. Returns: Nothing lol. """ self.model.fit(*args, **kwargs) def evaluate(self, *args, **kwargs): """evaluate. Literally use this as you'd use the normal keras evaluate. Args: Returns: The model's score, evaluated on whatever inputs you just fed it. """ return self.model.evaluate( *args, **kwargs) def predict(self, *args, **kwargs): """predict. Literally use this as you'd use the normal keras predict. Returns: a probability distribution over outputs, for each input. """ all_preds = self.model.predict(*args, **kwargs) return self.get_dist(all_preds) def get_dist(self, y_pred): """turns an output into a distribution. Literally use this as you'd use the normal keras predict. Args: y_pred: nn output Returns: a probability distribution over outputs, for each input. """ num_mix = self.num_mix output_dim = self.output_dim y_pred = tf.reshape(y_pred, [-1, (2 * num_mix * output_dim) + num_mix], name='reshape_ypreds') out_mu, out_sigma, out_pi = tf.split(y_pred, num_or_size_splits=[num_mix * output_dim, num_mix * output_dim, num_mix], axis=1, name='mdn_coef_split') cat = tfd.Categorical(logits=out_pi) component_splits = [output_dim] * num_mix mus = tf.split(out_mu, num_or_size_splits=component_splits, axis=1) sigs = tf.split(out_sigma, num_or_size_splits=component_splits, axis=1) coll = [tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale) for loc, scale in zip(mus, sigs)] return tfd.Mixture(cat=cat, components=coll)Methods
def compile(self, *args, loss=None, **kwargs)-
compile. Literally use this as you'd use the normal compile, but don't use your own loss functions, unless your really deep in this. Let the default one go
Args
loss- if you really wanna make your own loss function
Returns
Nothing lol
Expand source code
def compile(self, *args, loss=None, **kwargs): """compile. Literally use this as you'd use the normal compile, but don't use your own loss functions, unless your really deep in this. Let the default one go Args: loss: if you really wanna make your own loss function Returns: Nothing lol """ if loss is None: loss = mdn.get_mixture_loss_func(self.output_dim, self.num_mix) kwargs['loss'] = loss self.model.compile(*args, **kwargs) def elu_plus_one_plus_epsilon(self, x)-
ELU activation with a very small addition to help prevent NaN in loss.
Expand source code
def elu_plus_one_plus_epsilon(self, x): """ELU activation with a very small addition to help prevent NaN in loss.""" return tf.keras.backend.elu(x) + 1 + .00001 def evaluate(self, *args, **kwargs)-
evaluate. Literally use this as you'd use the normal keras evaluate. Args:
Returns
The model's score, evaluated on whatever inputs you just fed it.
Expand source code
def evaluate(self, *args, **kwargs): """evaluate. Literally use this as you'd use the normal keras evaluate. Args: Returns: The model's score, evaluated on whatever inputs you just fed it. """ return self.model.evaluate( *args, **kwargs) def fit(self, *args, **kwargs)-
fit. Literally use this as you'd use the normal keras fit.
Returns
Nothing lol.
Expand source code
def fit(self, *args, **kwargs): """fit. Literally use this as you'd use the normal keras fit. Returns: Nothing lol. """ self.model.fit(*args, **kwargs) def get_dist(self, y_pred)-
turns an output into a distribution. Literally use this as you'd use the normal keras predict.
Args
y_pred- nn output
Returns
a probability distribution over outputs, for each input.
Expand source code
def get_dist(self, y_pred): """turns an output into a distribution. Literally use this as you'd use the normal keras predict. Args: y_pred: nn output Returns: a probability distribution over outputs, for each input. """ num_mix = self.num_mix output_dim = self.output_dim y_pred = tf.reshape(y_pred, [-1, (2 * num_mix * output_dim) + num_mix], name='reshape_ypreds') out_mu, out_sigma, out_pi = tf.split(y_pred, num_or_size_splits=[num_mix * output_dim, num_mix * output_dim, num_mix], axis=1, name='mdn_coef_split') cat = tfd.Categorical(logits=out_pi) component_splits = [output_dim] * num_mix mus = tf.split(out_mu, num_or_size_splits=component_splits, axis=1) sigs = tf.split(out_sigma, num_or_size_splits=component_splits, axis=1) coll = [tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale) for loc, scale in zip(mus, sigs)] return tfd.Mixture(cat=cat, components=coll) def predict(self, *args, **kwargs)-
predict. Literally use this as you'd use the normal keras predict.
Returns
a probability distribution over outputs, for each input.
Expand source code
def predict(self, *args, **kwargs): """predict. Literally use this as you'd use the normal keras predict. Returns: a probability distribution over outputs, for each input. """ all_preds = self.model.predict(*args, **kwargs) return self.get_dist(all_preds)