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)