def GMM2Dslow(log_pis, mus, log_sigmas, corrs, clip_lo=-10, clip_hi=10):
# shapes
# pis: [..., GMM_c]
# mus: [..., GMM_c*state_dim]
# sigmas: [..., GMM_c*state_dim]
# corrs: [..., GMM_c]
GMM_c = log_pis.shape[-1]
mus_split = tf.split(mus, GMM_c, axis=-1)
sigmas = tf.exp(tf.clip_by_value(log_sigmas, clip_lo, clip_hi))
# Sigma = [s1^2 p*s1*s2 L = [s1 0
# p*s1*s2 s2^2 ] p*s2 sqrt(1-p^2)*s2]
sigmas_reshaped = tf.reshape(
sigmas,
[-1 if s.value is None else s.value
for s in sigmas.shape[:-1]] + [GMM_c.value, 2])
Ls = tf.stack(
[
(sigmas_reshaped *
tf.stack([tf.ones_like(corrs), corrs], -1)), # [s1, p*s2]
(sigmas_reshaped *
tf.stack([tf.zeros_like(corrs),
tf.sqrt(1 - corrs**2)], -1))
], # [0, sqrt(1-p^2)*s2]
axis=-1)
Ls_split = tf.unstack(Ls, axis=-3)
cat = distributions.Categorical(logits=log_pis)
dists = [
distributions.MultivariateNormalTriL(mu, L)
for mu, L in zip(mus_split, Ls_split)
]
return distributions.Mixture(cat, dists)
def _to_normal2d(output_batch) -> ds.MultivariateNormalTriL:
#lambada匿名函数
# 求x 和 y的平均值
x_mean = Lambda(lambda o: o[:, 0])(output_batch)
y_mean = Lambda(lambda o: o[:, 1])(output_batch)
# 求x 和 y的标准差，非负值
x_std = Lambda(lambda o: K.exp(o[:, 2]))(output_batch)
y_std = Lambda(lambda o: K.exp(o[:, 3]))(output_batch)
# 相关系数，范围为[-1,1]
cor = Lambda(lambda o: K.tanh(o[:, 4]))(output_batch)
loc = Concatenate()([
Lambda(lambda x_mean: K.expand_dims(x_mean, 1))(x_mean),
Lambda(lambda y_mean: K.expand_dims(y_mean, 1))(y_mean)
])
x_var = Lambda(lambda x_std: K.square(x_std))(x_std)
y_var = Lambda(lambda y_std: K.square(y_std))(y_std)
xy_cor = Multiply()([x_std, y_std, cor])
cov = Lambda(lambda inputs: K.stack(inputs, axis=0))(
[x_var, xy_cor, xy_cor, y_var])
cov = Lambda(lambda cov: K.permute_dimensions(cov, (1, 0)))(cov)
cov = Reshape((2, 2))(cov)
scale_tril = Lambda(lambda cov: tf.cholesky(cov))(cov)
mvn = ds.MultivariateNormalTriL(loc, scale_tril)
return mvn
def test_blobs():
from matplotlib import pyplot as plt
tf.enable_eager_execution()
import numpy as np
import tensorflow.contrib.distributions as tfd
_means = [-0.5, 0, 0.5]
means = tf.ones((3, 1, 2), dtype=tf.float32) * np.array(_means).reshape(
(3, 1, 1))
means = tf.concat([means, means, means[::-1, ...]], axis=1)
means = tf.reshape(means, (-1, 2))
var_ = 0.1
rho = 0.5
cov = [[var_, rho * var_], [rho * var_, var_]]
scale = tf.cholesky(cov)
scale = tf.stack([scale] * 3, axis=0)
scale = tf.stack([scale] * 3, axis=0)
scale = tf.reshape(scale, (-1, 2, 2))
mvn = tfd.MultivariateNormalTriL(loc=means, scale_tril=scale)
h = 100
w = 100
y_t = tf.tile(tf.reshape(tf.linspace(-1.0, 1.0, h), [h, 1]), [1, w])
x_t = tf.tile(tf.reshape(tf.linspace(-1.0, 1.0, w), [1, w]), [h, 1])
y_t = tf.expand_dims(y_t, axis=-1)
x_t = tf.expand_dims(x_t, axis=-1)
meshgrid = tf.concat([y_t, x_t], axis=-1)
meshgrid = tf.expand_dims(meshgrid, 0)
meshgrid = tf.expand_dims(meshgrid, 3) # 1, h, w, 1, 2
blob = mvn.prob(meshgrid)
blob = tf.reshape(blob, (100, 100, 3, 3))
blob = tf.transpose(blob, perm=[2, 0, 1, 3])
norm_const = np.sum(blob, axis=(1, 2), keepdims=True)
mu, L = nn.probs_to_mu_L(blob / norm_const, 1, inv=False)
bn, h, w, nk = blob.get_shape().as_list()
estimated_blob = nn.tf_hm(h, w, mu, L)
fig, ax = plt.subplots(2, 3, figsize=(9, 6))
for b in range(len(_means)):
ax[0, b].imshow(np.squeeze(blob[b, ...]))
ax[0, b].set_title("target_blobs")
ax[0, b].set_axis_off()
for b in range(len(_means)):
ax[1, b].imshow(np.squeeze(estimated_blob[b, ...]))
ax[1, b].set_title("estimated_blobs")
ax[1, b].set_axis_off()
def MultivariateGaussianTril(mu, scale):
'''
Create multivariate Gaussian distribution from mu and lower trianglar
covariance matrix.
Parameters
----------
mu : list
mean parameters
scale : list
scale of multivariate gaussian. covariance = scale @ scale.T
Returns
-------
Multivariate Gaussian distribution
'''
scale_tril = tfd.fill_triangular(scale, upper=False)
mvn = tfd.MultivariateNormalTriL(loc=mu, scale_tril=scale_tril)
return mvn
def get_mixture_model(self):
"""
ds.Mixture in TensorFlow requires a Categorical dist. to determine which individual dist. is
used for generating a sample, 'components' is a list of different classes defined from
tf.contrib.distributions
"""
prob = 1. / self.num_gaussians
probs = [prob for i in range(self.num_gaussians)]
mus = self.get_mus()
scales = self.get_scale_matrices()
gaussians = [
ds.MultivariateNormalTriL(loc=mus[i], scale_tril=scales[i])
for i in range(self.num_gaussians)
]
mixture = ds.Mixture(cat=ds.Categorical(probs=probs),
components=gaussians)
return mixture
def _to_normal2d(output_batch) -> ds.MultivariateNormalTriL:
"""
:param output_batch: (n_samples, 5)
:return
"""
# mean of x and y
x_mean = Lambda(lambda o: o[:, 0])(output_batch)
y_mean = Lambda(lambda o: o[:, 1])(output_batch)
# std of x and y
# std is must be 0 or positive
x_std = Lambda(lambda o: K.exp(o[:, 2]))(output_batch)
y_std = Lambda(lambda o: K.exp(o[:, 3]))(output_batch)
# correlation coefficient
# correlation coefficient range is [-1, 1]
cor = Lambda(lambda o: K.tanh(o[:, 4]))(output_batch)
loc = Concatenate()([
Lambda(lambda x_mean: K.expand_dims(x_mean, 1))(x_mean),
Lambda(lambda y_mean: K.expand_dims(y_mean, 1))(y_mean)
])
x_var = Lambda(lambda x_std: K.square(x_std))(x_std)
y_var = Lambda(lambda y_std: K.square(y_std))(y_std)
xy_cor = Multiply()([x_std, y_std, cor])
cov = Lambda(lambda inputs: K.stack(inputs, axis=0))(
[x_var, xy_cor, xy_cor, y_var])
cov = Lambda(lambda cov: K.permute_dimensions(cov, (1, 0)))(cov)
cov = Reshape((2, 2))(cov)
scale_tril = Lambda(lambda cov: tf.cholesky(cov))(cov)
mvn = ds.MultivariateNormalTriL(loc, scale_tril)
return mvn
def __init__(self, *args, **kwargs):
super(AffineDiagNormal, self).__init__()
self._factory = lambda loc, scale_tril: tfd.MultivariateNormalTriL(
loc, scale_tril, *args, **kwargs)
def __init__(self, sess, ob_space, ac_space, nbatch, nsteps, reuse=False,
name='policy', args=None): #pylint: disable=W0613
policy_variance_state_dependent = args.policy_variance_state_dependent
ac_fn = args.ac_fn
hidden_sizes = args.hidden_sizes
num_sharing_layers = args.num_sharing_layers
num_layers = args.num_layers
assert ac_fn in ['tanh', 'sigmoid', 'relu']
if isinstance(hidden_sizes, int):
assert num_layers is not None
hidden_sizes = [hidden_sizes] * num_layers
if num_layers is None:
num_layers = len(hidden_sizes)
assert num_layers == len(hidden_sizes)
# print(f'Policy hidden_sizes:{hidden_sizes}')
self.pdtype = make_pdtype(ac_space)
with tf.variable_scope(name, reuse=reuse):
X, processed_x = observation_input(ob_space, nbatch)
activ = getattr( tf.nn, ac_fn )
processed_x = tf.layers.flatten(processed_x)
# --- share layers
for ind_layer in range(num_sharing_layers):
processed_x = activ( fc(processed_x, f'share_fc{ind_layer}', nh=hidden_sizes[ind_layer], init_scale=np.sqrt(2)) )
# --- policy
pi_h = processed_x
for ind_layer in range( num_sharing_layers, num_layers ):
pi_h = activ(fc(pi_h, f'pi_fc{ind_layer}', nh=hidden_sizes[ind_layer], init_scale=np.sqrt(2)))
from gym import spaces
params_addtional = {}
if policy_variance_state_dependent and isinstance( ac_space, spaces.Box ):
latent_logstd = processed_x
for ind_layer in range(num_sharing_layers, num_layers):
latent_logstd = activ(fc(latent_logstd, f'logstd_fc{ind_layer}', nh=hidden_sizes[ind_layer], init_scale=np.sqrt(2)))
params_addtional['latent_logstd'] = latent_logstd
self.pd, self.pi = self.pdtype.pdfromlatent(pi_h, init_scale=0.01, logstd_initial=args.logstd, **params_addtional)
# --- value function
vf_h = processed_x
for ind_layer in range( num_sharing_layers, num_layers ):
vf_h = activ(fc(vf_h, f'vf_fc{ind_layer}', nh=hidden_sizes[ind_layer], init_scale=np.sqrt(2)))
vf = fc(vf_h, 'vf', 1)[:,0]
a_sample = self.pd.sample()
neglogp_sample = self.pd.neglogp(a_sample)
self.initial_state = None
# --- predict function
# use placeholder
# use stochastic action
# use deterministic action
if args.coef_predict_task > 0:
import tensorflow.contrib.distributions as dists
assert isinstance( ac_space, Box ), 'Only Implement for Box action space'
A_type = tf.placeholder_with_default('pl', dtype=tf.string)
A_pl = self.pdtype.sample_placeholder([None])
self.A = A_pl
self.A_type = A_type
A_input_1 = U.switch( tf.equal( A_type, 'det' ), self.pd.mode(), a_sample )
A_input = U.switch( tf.equal( A_type, 'pl' ), A_pl,A_input_1)
predict_h = tf.concat( (processed_x, A_input))
for ind_layer in range(num_sharing_layers, num_layers):
predict_h = activ(fc(predict_h, f'predict_fc{ind_layer}', nh=hidden_sizes[ind_layer], init_scale=np.sqrt(2)))
predict_mean = fc(predict_h, f'predict_fc{ind_layer}', nh=ob_space.shape[0], init_scale=np.sqrt(2))
predict_cov_init_value = np.identity( shape=ob_space.shape )
predict_cov = tf.get_variable( name='predict_cov', shape=predict_cov_init_value, initializer=tf.constant_initializer(predict_cov_init_value) )
predict_dist = dists.MultivariateNormalTriL( predict_mean, predict_cov )
self.predict_dist = predict_dist
scope_model = tf.get_variable_scope().name
self.variables_all = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope_model)
self.variables_trainable = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope_model)
#--- set logstd
# if isinstance( ac_space, Box ):
# if not policy_variance_state_dependent:
# logstd_pl, _ = observation_input( ac_space, batch_size=1, name='ac' )
# assign_logstd = tf.assign( self.pdtype.logstd, logstd_pl )
# set_logstd_entity = U.function([logstd_pl], assign_logstd)
# def set_logstd(logstd_new):
# # if isinstance( logstd_new, float ):
# # logstd_new = [[logstd_new] * ac_space.shape[0]]
# set_logstd_entity(logstd_new)
# self.set_logstd = set_logstd
# self.get_logstd = U.function([], self.pdtype.logstd)
def step(ob, *_args, **_kwargs):
a, v, neglogp = sess.run([a_sample, vf, neglogp_sample], {X:ob})
return a, v, self.initial_state, neglogp
def value(ob, *_args, **_kwargs):
return sess.run(vf, {X:ob})
def step_policyflat(ob, *_args, **_kwargs):
a, v, neglogp, polciyflat = sess.run([a_sample, vf, neglogp_sample, self.pd.flatparam()], {X:ob}) #TODO: TEST flat for discrete action space
return a, v, self.initial_state, neglogp, polciyflat
def step_test(ob, *_args, **_kwargs):
a = sess.run([self.pd.mode()], {X:ob})
return a
self.X = X
self.vf = vf
self.step = step
self.step_policyflat = step_policyflat
self.value = value
self.step_test = step_test
def tf_hm(h, w, mu, L, order="exp"):
"""
Returns Gaussian densitiy function based on μ and L for each batch index and part
L is the cholesky decomposition of the covariance matrix : Σ = L L^T
Parameters
----------
h : int
heigh ot output map
w : int
width of output map
mu : tensor
mean of gaussian part and batch item. Shape [b, p, 2]. Mean in range [-1, 1] with respect to height and width
L : tensor
cholesky decomposition of covariance matrix for each batch item and part. Shape [b, p, 2, 2]
order:
Returns
-------
density : tensor
gaussian blob for each part and batch idx. Shape [b, h, w, p]
Examples
--------
from matplotlib import pyplot as plt
tf.enable_eager_execution()
import numpy as np
import tensorflow as tf
import tensorflow.contrib.distributions as tfd
# create Target Blobs
_means = [-0.5, 0, 0.5]
means = tf.ones((3, 1, 2), dtype=tf.float32) * np.array(_means).reshape((3, 1, 1))
means = tf.concat([means, means, means[::-1, ...]], axis=1)
means = tf.reshape(means, (-1, 2))
var_ = 0.1
rho = 0.5
cov = [[var_, rho * var_],
[rho * var_, var_]]
scale = tf.cholesky(cov)
scale = tf.stack([scale] * 3, axis=0)
scale = tf.stack([scale] * 3, axis=0)
scale = tf.reshape(scale, (-1, 2, 2))
mvn = tfd.MultivariateNormalTriL(
loc=means,
scale_tril=scale)
h = 100
w = 100
y_t = tf.tile(tf.reshape(tf.linspace(-1., 1., h), [h, 1]), [1, w])
x_t = tf.tile(tf.reshape(tf.linspace(-1., 1., w), [1, w]), [h, 1])
y_t = tf.expand_dims(y_t, axis=-1)
x_t = tf.expand_dims(x_t, axis=-1)
meshgrid = tf.concat([y_t, x_t], axis=-1)
meshgrid = tf.expand_dims(meshgrid, 0)
meshgrid = tf.expand_dims(meshgrid, 3) # 1, h, w, 1, 2
blob = mvn.prob(meshgrid)
blob = tf.reshape(blob, (100, 100, 3, 3))
blob = tf.transpose(blob, perm=[2, 0, 1, 3])
# Estimate mean and L
norm_const = np.sum(blob, axis=(1, 2), keepdims=True)
mu, L = nn.probs_to_mu_L(blob / norm_const, 1, inv=False)
bn, h, w, nk = blob.get_shape().as_list()
# Estimate blob based on mu and L
estimated_blob = nn.tf_hm(h, w, mu, L)
# plot
fig, ax = plt.subplots(2, 3, figsize=(9, 6))
for b in range(len(_means)):
ax[0, b].imshow(np.squeeze(blob[b, ...]))
ax[0, b].set_title("target_blobs")
ax[0, b].set_axis_off()
for b in range(len(_means)):
ax[1, b].imshow(np.squeeze(estimated_blob[b, ...]))
ax[1, b].set_title("estimated_blobs")
ax[1, b].set_axis_off()
"""
assert len(mu.get_shape().as_list()) == 3
assert len(L.get_shape().as_list()) == 4
assert mu.get_shape().as_list()[-1] == 2
assert L.get_shape().as_list()[-1] == 2
assert L.get_shape().as_list()[-2] == 2
b, p, _ = mu.get_shape().as_list()
mu = tf.reshape(mu, (b * p, 2))
L = tf.reshape(L, (b * p, 2, 2))
mvn = tfd.MultivariateNormalTriL(loc=mu, scale_tril=L)
y_t = tf.tile(tf.reshape(tf.linspace(-1.0, 1.0, h), [h, 1]), [1, w])
x_t = tf.tile(tf.reshape(tf.linspace(-1.0, 1.0, w), [1, w]), [h, 1])
y_t = tf.expand_dims(y_t, axis=-1)
x_t = tf.expand_dims(x_t, axis=-1)
meshgrid = tf.concat([y_t, x_t], axis=-1)
meshgrid = tf.expand_dims(meshgrid, 0)
meshgrid = tf.expand_dims(meshgrid, 3) # 1, h, w, 1, 2
probs = mvn.prob(meshgrid)
probs = tf.reshape(probs, (h, w, b, p))
probs = tf.transpose(probs, perm=[2, 0, 1,
3]) # move part axis to the back
return probs
def test_probs_to_mu_sigma():
# TODO: right now, the test is only visual debugging
# We would need a numeric criterion to test if the calculation is correct
from matplotlib import pyplot as plt
tf.enable_eager_execution()
import numpy as np
import tensorflow.contrib.distributions as tfd
_means = [-0.5, 0, 0.5]
means = tf.ones((3, 1, 2), dtype=tf.float32) * np.array(_means).reshape(
(3, 1, 1))
means = tf.concat([means, means, means[::-1, ...]], axis=1)
means = tf.reshape(means, (-1, 2))
var_ = 0.1
rho = 0.0
cov = [[var_, rho * var_], [rho * var_, var_]]
scale = tf.cholesky(cov)
scale = tf.stack([scale] * 3, axis=0)
scale = tf.stack([scale] * 3, axis=0)
scale = tf.reshape(scale, (-1, 2, 2))
mvn = tfd.MultivariateNormalTriL(loc=means, scale_tril=scale)
h = 100
w = 100
y_t = tf.tile(tf.reshape(tf.linspace(-1.0, 1.0, h), [h, 1]), [1, w])
x_t = tf.tile(tf.reshape(tf.linspace(-1.0, 1.0, w), [1, w]), [h, 1])
y_t = tf.expand_dims(y_t, axis=-1)
x_t = tf.expand_dims(x_t, axis=-1)
meshgrid = tf.concat([y_t, x_t], axis=-1)
meshgrid = tf.expand_dims(meshgrid, 0)
meshgrid = tf.expand_dims(meshgrid, 3) # 1, h, w, 1, 2
blob = mvn.prob(meshgrid)
blob = tf.reshape(blob, (100, 100, 3, 3))
blob = tf.transpose(blob, perm=[2, 0, 1, 3])
norm_const = np.sum(blob, axis=(1, 2), keepdims=True)
mu, sigma = nn.probs_to_mu_sigma(blob / norm_const)
# norm_const2 = np.sum(blob, axis=(1, 2), keepdims=False)
# mu2, sigma2 = nn.probs_to_mu_sigma(blob, 1 / norm_const2)
#
# assert np.allclose(mu, mu2, rtol=1e-4, atol=1e-4)
# assert np.allclose(sigma, sigma2, rtol=1e-4, atol=1e-4)
L = tf.cholesky(sigma)
bn, h, w, nk = blob.get_shape().as_list()
estimated_blob = nn.tf_hm(h, w, mu, L)
fig, ax = plt.subplots(2, 3, figsize=(9, 6))
for b in range(len(_means)):
ax[0, b].imshow(np.squeeze(blob[b, ...]))
ax[0, b].set_title("target_blobs")
ax[0, b].set_axis_off()
for b in range(len(_means)):
ax[1, b].imshow(np.squeeze(estimated_blob[b, ...]))
ax[1, b].set_title("estimated_blobs")
ax[1, b].set_axis_off()
plt.show()
def do_filter(self, estimated_state, estimated_state_covariance,
predicted_observation, predicted_observation_covariance,
observation, observation_model, observation_noise):
"""Convenience function for scoring predictions.
Scores a prediction against an observation, and computes the updated
posterior over states.
Shapes given below for arguments are for single-model Kalman filtering
(e.g. KalmanFilter). For ensembles, prior_state and prior_state_var are
same-length tuples of values corresponding to each model.
Args:
estimated_state: A prior mean over states [batch size x state dimension]
estimated_state_covariance: Covariance of state prior [batch size x D x
D], with D depending on the Kalman filter implementation (typically
the state dimension).
predicted_observation: A prediction for the observed value, such as that
returned by observed_from_state. A [batch size x num features] Tensor.
predicted_observation_covariance: A covariance matrix corresponding to
`predicted_observation`, a [batch size x num features x num features]
Tensor.
observation: The observed value corresponding to the predictions
given [batch size x observation dimension]
observation_model: The [batch size x observation dimension x model state
dimension] Tensor indicating how a particular state is mapped to
(pre-noise) observations for each part of the batch.
observation_noise: A [batch size x observation dimension x observation
dimension] Tensor or [observation dimension x observation dimension]
Tensor with covariance matrices to use for each part of the batch (a
two-dimensional input will be broadcast).
Returns:
posterior_state, posterior_state_var: Posterior mean and
covariance, updated versions of prior_state and
prior_state_var.
log_prediction_prob: Log probability of the observations under
the priors, suitable for optimization (should be maximized).
"""
symmetrized_observation_covariance = 0.5 * (
predicted_observation_covariance +
array_ops.matrix_transpose(predicted_observation_covariance))
instability_message = (
"This may occur due to numerically unstable filtering when there is "
"a large difference in posterior variances, or when inferences are "
"near-deterministic. Considering tuning the "
"'filtering_maximum_posterior_variance_ratio' or "
"'filtering_minimum_posterior_variance' parameters in your "
"StateSpaceModelConfiguration, or tuning the transition matrix.")
symmetrized_observation_covariance = numerics.verify_tensor_all_finite(
symmetrized_observation_covariance,
"Predicted observation covariance was not finite. {}".format(
instability_message))
diag = array_ops.matrix_diag_part(symmetrized_observation_covariance)
min_diag = math_ops.reduce_min(diag)
non_negative_assert = control_flow_ops.Assert(
min_diag >= 0.,
[("The predicted observation covariance "
"has a negative diagonal entry. {}").format(instability_message),
min_diag])
with ops.control_dependencies([non_negative_assert]):
observation_covariance_cholesky = linalg_ops.cholesky(
symmetrized_observation_covariance)
log_prediction_prob = distributions.MultivariateNormalTriL(
predicted_observation,
observation_covariance_cholesky).log_prob(observation)
(posterior_state,
posterior_state_var) = self.posterior_from_prior_state(
prior_state=estimated_state,
prior_state_var=estimated_state_covariance,
observation=observation,
observation_model=observation_model,
predicted_observations=(predicted_observation,
predicted_observation_covariance),
observation_noise=observation_noise)
return (posterior_state, posterior_state_var, log_prediction_prob)
if args.SINGLEVAR == 0:
scale_px = activations[args.SFPX](bias_variable(
[args.NX], value=[float(args.init_sigma)] * args.NZ, name='scale'))
loc = bias_variable([args.NX], value=0., name='loc')
if args.ZEROMEANPX == 1:
loc = tf.zeros_like(loc)
pdfError = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_px)
if args.SINGLEVAR == 2:
scale_px = activations[args.SFPX](bias_variable(
[args.NX * (args.NX + 1) // 2], value=-1., name='scale'))
loc = bias_variable([args.NX], value=0., name='loc')
if args.ZEROMEANPX == 1:
loc = tf.zeros_like(loc)
pdfError = tfd.MultivariateNormalTriL(
loc=loc, scale_tril=tfd.fill_triangular(scale_px))
###########################################################################################
with tf.variable_scope('prior', reuse=tf.AUTO_REUSE):
prior = create_gmm_1(args.NZ,
args.PRIOR_NK,
'prior',
scale_act=activations[args.SF],
zero_mean=args.ZM > 0)
# prior = tfd.Independent( tfd.Laplace( loc = [ [0.]*args.NZ , ], scale=[ [1.]*args.NZ ,] ))
############################################################################################
config = tf.ConfigProto(device_count={'GPU': args.GPU})
config.log_device_placement = False
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
# sess = tf.Session()
def _mdn_model_fn(features, labels, hidden_units, n_mixture, diag,
feature_columns, label_columns, optimizer, activation_fn,
normalizer_fn, dropout, mode):
# Check for training mode
is_training = mode == tf.estimator.ModeKeys.TRAIN
label_dimension = len(label_columns)
# Extracts the features
input_layer = tf.feature_column.input_layer(
features=features, feature_columns=feature_columns)
# Builds the neural network
#net = deep_mlp(input_layer, hidden_units, activation_fn, normalizer_fn,
# dropout, is_training)
net = slim.fully_connected(input_layer, 512, activation_fn=tf.nn.elu)
net = slim.fully_connected(net, 512, activation_fn=tf.nn.elu)
# Size of the covariance matrix
if diag == True:
size_sigma = label_dimension
else:
size_sigma = (label_dimension * (label_dimension + 1) // 2)
# Create mixture components from network output
out_mu = slim.fully_connected(net,
label_dimension * n_mixture,
activation_fn=None)
out_mu = tf.reshape(out_mu, (-1, n_mixture, label_dimension))
out_sigma = slim.fully_connected(net,
size_sigma * n_mixture,
activation_fn=None)
out_sigma = tf.reshape(out_sigma, (-1, n_mixture, size_sigma))
out_p = slim.fully_connected(net, n_mixture, activation_fn=None)
out_p = tf.nn.softmax(tf.reshape(out_p, (-1, n_mixture)))
if diag == True:
sigma_mat = tf.nn.softplus(out_sigma) + 1e-4
gmm = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(probs=out_p),
components_distribution=tfd.MultivariateNormalDiag(
loc=out_mu, scale_diag=sigma_mat))
else:
sigma_mat = tfd.matrix_diag_transform(tfd.fill_triangular(out_sigma),
transform=tf.nn.softplus)
gmm = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(probs=out_p),
components_distribution=tfd.MultivariateNormalTriL(
loc=out_mu, scale_tril=sigma_mat))
predictions = {'mu': out_mu, 'sigma': sigma_mat, 'p': out_p}
if mode == tf.estimator.ModeKeys.PREDICT:
# Add functionality to sample directly from the model
y = gmm.sample()
samples = {}
for i, k in enumerate(label_columns):
samples[k.name] = y[:, -i - 1]
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=predictions,
export_outputs={
'pdf':
tf.estimator.export.PredictOutput(predictions),
'samples':
tf.estimator.export.PredictOutput(samples),
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
tf.estimator.export.PredictOutput(samples)
})
label_layer = tf.feature_column.input_layer(features=features,
feature_columns=label_columns)
# Compute and register loss function
loss = -tf.reduce_mean(gmm.log_prob(label_layer), axis=0)
tf.losses.add_loss(loss)
total_loss = tf.losses.get_total_loss()
train_op = None
eval_metric_ops = None
# Define optimizer
if mode == tf.estimator.ModeKeys.TRAIN:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer(learning_rate=0.00001).minimize(
loss=total_loss, global_step=tf.train.get_global_step())
tf.summary.scalar('loss', loss)
elif mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {"log_p": loss}
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=total_loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)
