from __future__ import absolute_import, division, print_function, unicode_literals
import numpy
import keras
from keras import backend as K
from keras.layers.merge import Concatenate
from keras.layers import Lambda
from keras.layers.core import Reshape
def split(start, stop):
return Lambda(lambda x: x[:, start:stop], output_shape=(None, stop-start))
def split_mixture_of_gaussians(x, n_components):
pi = split(0, n_components)(x)
mu = split(n_components, 2*n_components)(x)
log_sig = split(2*n_components, 3*n_components)(x)
return pi, mu, log_sig
def log_norm_pdf(x, mu, log_sig):
z = (x - mu) / (K.exp(K.clip(log_sig, -40, 40))) #TODO: get rid of this clipping
return -(0.5)*K.log(2*numpy.pi) - log_sig - 0.5*((z)**2)
def mix_gaussian_loss(x, mu, log_sig, w):
'''
Combine the mixture of gaussian distribution and the loss into a single function
so that we can do the log sum exp trick for numerical stability...
'''
if K.backend() == "tensorflow":
x.set_shape([None, 1])
gauss = log_norm_pdf(K.repeat_elements(x=x, rep=mu.shape[1], axis=1), mu, log_sig)
# TODO: get rid of clipping.
gauss = K.clip(gauss, -40, 40)
max_gauss = K.maximum((0.), K.max(gauss))
# log sum exp trick...
gauss = gauss - max_gauss
out = K.sum(w * K.exp(gauss), axis=1)
loss = K.mean(-K.log(out) + max_gauss)
return loss
def mixture_of_gaussian_output(x, n_components):
mu = keras.layers.Dense(n_components, activation='linear')(x)
log_sig = keras.layers.Dense(n_components, activation='linear')(x)
pi = keras.layers.Dense(n_components, activation='softmax')(x)
return Concatenate(axis=1)([pi, mu, log_sig])
def mixture_of_gaussian_loss(y_true, y_pred, n_components):
pi, mu, log_sig = split_mixture_of_gaussians(y_pred, n_components)
return mix_gaussian_loss(y_true, mu, log_sig, pi)
def mixture_gaussian(n_components):
'''
Build a mixture of gaussian output and loss function that may be used in a keras graph.
'''
def output(x):
return mixture_of_gaussian_output(x, n_components)
def keras_loss(y, x):
return mixture_of_gaussian_loss(y, x, n_components)
return output, keras_loss
