def _build(self, inputs): locs = self._modules['locs'](inputs) log_scales = self._modules['scales'](inputs) logits = self._modules['logits'](inputs) scales = tf.nn.softplus(log_scales + softplus_inverse(1.0)) locs = tf.reshape(locs, [-1, self.k, self.ndim]) scales = tf.reshape(scales, [-1, self.k, self.ndim]) logits = tf.reshape(logits, [-1, self.k]) # reshape so that the first dim is the mixture, because we are doing to unstack them # also swap the batch size and the ones that come from the steps of this run # (K x N x D) mix_first_locs = tf.transpose(locs, [1, 0, 2]) mix_first_scales = tf.transpose(scales, [1, 0, 2]) outs = {'locs': locs, 'scales': scales, 'logits': logits} outs['flattened'] = flatten_mdn(logits, locs, scales, self.FLAGS) cat = tfd.Categorical(logits=logits) components = [] eval_components = [] for loc, scale in zip(tf.unstack(mix_first_locs), tf.unstack(mix_first_scales)): normal = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale) components.append(normal) eval_normal = tfd.MultivariateNormalDiag(loc=loc[..., :2], scale_diag=scale[..., :2]) eval_components.append(eval_normal) mixture = tfd.Mixture(cat=cat, components=components) eval_cat = tfd.Categorical(logits=logits) eval_mixture = tfd.Mixture(cat=eval_cat, components=eval_components) outs['mixture'] = mixture outs['eval_mixture'] = eval_mixture return outs
def get_pdf(param_vec, vehicle_type): # see https://ericmjl.github.io/blog/2019/5/29/reasoning-about-shapes-and-probability-distributions/ # for info on shapes if vehicle_type == 'other_vehicle': alpha, mus, sigmas = slice_pvector(param_vec, vehicle_type) # Unpack parameter vectors mvn = tfd.MixtureSameFamily( mixture_distribution=tfd.Categorical(probs=alpha), components_distribution=tfd.Normal( loc=mus, scale=sigmas)) if vehicle_type == 'merge_vehicle': alphas, mus_long, sigmas_long, mus_lat, \ sigmas_lat, rhos = slice_pvector(param_vec, vehicle_type) cov = get_CovMatrix(rhos, sigmas_long, sigmas_lat) mus = tf.stack([mus_long, mus_lat], axis=3, name='mus') mvn = tfd.MixtureSameFamily( mixture_distribution=tfd.Categorical( probs=alphas), components_distribution=tfd.MultivariateNormalTriL( loc=mus, scale_tril=tf.linalg.cholesky(cov), name='MultivariateNormalTriL')) # print('mus shape: ', mus.shape) return mvn
def conditional_distribution(self, x, slice_in, slice_out): marginal_in = ds.MixtureSameFamily( mixture_distribution=ds.Categorical(probs=self._priors), components_distribution=ds.MultivariateNormalFullCovariance( loc=self._locs[:, slice_in], covariance_matrix=self._covs[:, slice_in, slice_in])) p_k_in = ds.Categorical(logits=tf_utils.log_normalize( marginal_in.components_distribution.log_prob(x[:, None]) + marginal_in.mixture_distribution.logits[None], axis=1)) sigma_in_out = self._covs[:, slice_in, slice_out] inv_sigma_in_in = tf.linalg.inv(self._covs[:, slice_in, slice_in]) inv_sigma_out_in = tf.matmul(sigma_in_out, inv_sigma_in_in, transpose_a=True) A = inv_sigma_out_in b = self._locs[:, slice_out] - tf.matmul( inv_sigma_out_in, self._locs[:, slice_in, None])[:, :, 0] cov_est = (self._covs[:, slice_out, slice_out] - tf.matmul(inv_sigma_out_in, sigma_in_out)) ys = tf.einsum('aij,bj->abi', A, x) + b[:, None] p_out_in_k = ds.MultivariateNormalFullCovariance( tf.transpose(ys, perm=(1, 0, 2)), cov_est) return ds.MixtureSameFamily(mixture_distribution=p_k_in, components_distribution=p_out_in_k)
def call(self, x, evaluation=False, epsilon=0, temp=1): self.logits = logits = self._layers(x) if evaluation: temp = tf.where(temp == 0, 1e-9, temp) logits = logits / temp dist = tfd.Categorical(logits) action = dist.sample() else: if self._epsilon_scaled_logits and \ (isinstance(epsilon, tf.Tensor) or epsilon): self.logits = epsilon_scaled_logits(logits, epsilon, temp) else: self.logits = logits / temp dist = tfd.Categorical(self.logits) action = dist.sample() if not self._epsilon_scaled_logits and \ (isinstance(epsilon, tf.Tensor) or epsilon): action = epsilon_greedy(action, epsilon, True, self.action_dim) self._dist = dist self._action = action return action
def mix(gamma, eta, loc, scale, neg_inf, n): return tfd.Mixture( cat=tfd.Categorical(probs=tf.stack([gamma, 1 - gamma], axis=-1)), components=[ tfd.Sample(tfd.Normal(np.float64(neg_inf), 1e-5), sample_shape=n), tfd.Sample(tfd.MixtureSameFamily( mixture_distribution=tfd.Categorical(probs=eta), components_distribution=tfd.Normal(loc=loc, scale=scale)), sample_shape=n) ])
def mix(gamma, eta, loc, scale, neg_inf): _gamma = gamma[..., tf.newaxis] # FIXME: Possible to use tfd.Blockwise? return tfd.Mixture( cat=tfd.Categorical(probs=tf.concat([_gamma, 1 - _gamma], axis=-1)), components=[ tfd.Deterministic(np.float64(neg_inf)), tfd.MixtureSameFamily( mixture_distribution=tfd.Categorical(probs=eta), components_distribution=tfd.Normal(loc=loc, scale=scale)), ])
def __call__(self): """Get the distribution object from the backend""" if get_backend() == 'pytorch': import torch.distributions as tod raise NotImplementedError else: from tensorflow_probability import distributions as tfd return tfd.HiddenMarkovModel( initial_distribution=tfd.Categorical(self['initial']), transition_distribution=tfd.Categorical(self['transition']), observation_distribution=self['observation'], num_steps=self['steps'])
def train_HMM(self): # Define variable to represent the unknown log rates. _trainable_log_rates = tf.Variable( np.log(np.mean(self.observed_counts)) + tf.random.normal([self.num_states]), name='log_rates') self.hmm = tfd.HiddenMarkovModel( initial_distribution=tfd.Categorical(logits=initial_state_logits), transition_distribution=tfd.Categorical( probs=self._transition_probs), observation_distribution=tfd.Poisson( log_rate=_trainable_log_rates), num_steps=len(observed_counts))
def call(self, inputs): prob, y = inputs uni_prob = K.ones_like(prob)/tf.cast(K.shape(prob)[-1], tf.float32) #variational distribution var_dist = tfp.Categorical(prob) #prior uniform distribution pri_dist = tfp.Categorical(uni_prob) ent_loss = K.mean(var_dist.entropy()) kl_loss = K.mean(tfp.kl_divergence(var_dist, pri_dist)) cent_loss = K.mean(K.categorical_crossentropy(y, prob)) self.add_loss( self.rho*cent_loss + \ (1-self.rho)*(self.lamb_kl*kl_loss+self.lamb_ent*ent_loss) ) return kl_loss
def mdn_loss_func(y_true, y_pred): # Reshape inputs in case this is used in a TimeDistribued layer y_pred = tf.reshape(y_pred, [-1, (2 * num_mixes * output_dim) + num_mixes], name='reshape_ypreds') y_true = tf.reshape(y_true, [-1, output_dim], name='reshape_ytrue') # Split the inputs into paramaters out_mu, out_sigma, out_pi = tf.split(y_pred, num_or_size_splits=[ num_mixes * output_dim, num_mixes * output_dim, num_mixes ], axis=-1, name='mdn_coef_split') # Construct the mixture models cat = tfd.Categorical(logits=out_pi) component_splits = [output_dim] * num_mixes 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) ] mixture = tfd.Mixture(cat=cat, components=coll) loss = mixture.log_prob(y_true) loss = tf.negative(loss) loss = tf.reduce_mean(loss) return loss
def generate_grm_data(n_sample, n_factor, n_item, nu, ld, rho, dtype = tf.float64): if (n_item % n_factor) != 0: n_item = n_factor * (n_item // n_factor) item_per_factor = (n_item // n_factor) n_category = len(nu) + 1 intercept = tf.tile(tf.constant([nu], dtype = dtype), multiples = [n_item, 1]) loading = np.zeros((n_item, n_factor)) for i in range(n_factor): for j in range(i * item_per_factor, (i + 1) * item_per_factor): loading[j, i] = ld loading = tf.constant(loading, dtype = dtype) if rho is None: cor = tf.eye(n_factor, dtype = dtype) else: unit = tf.ones((n_factor, 1), dtype = dtype) identity = tf.eye(n_factor, dtype = dtype) cor = rho * (unit @ tf.transpose(unit)) + (1 - rho) * identity dist_eta = tfd.MultivariateNormalTriL( loc = tf.zeros(n_factor, dtype = dtype), scale_tril = tf.linalg.cholesky(cor)) eta = dist_eta.sample(n_sample) c, d = create_cd(n_category, dtype) probs = grm_irf(eta, intercept, loading, c, d) x = tfd.Categorical(probs=probs, dtype=dtype).sample() return x
def mix(n, eta, loc, scale, name): return tfd.Sample( tfd.MixtureSameFamily( mixture_distribution=tfd.Categorical(probs=eta), components_distribution=tfd.Normal(loc=loc, scale=scale), name=name), sample_shape=n)
def call(self, inputs, **kwargs): latent = self._base_model(inputs) logits = self._policy_fn(latent) deterministic = kwargs.get("deterministic", False) if self._discrete: pi = tfd.Categorical(logits) if deterministic: actions = pi.mode() else: actions = pi.sample() logpacs = tf.nn.log_softmax(logits) else: means, logstds = tf.split(logits, 2, axis=-1) logstds = tf.clip_by_value(logstds, MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT) pi = tfd.MultivariateNormalDiag(means, tf.exp(logstds)) if deterministic: actions = pi.mean() else: actions = pi.sample() logpacs = pi.log_prob(actions) # Adjust loglikelihoods for squashed actions logpacs -= tf.reduce_sum( 2 * (np.log(2) - actions - tf.nn.softplus(-2 * actions)), axis=1 ) actions = tf.math.tanh(actions) actions = tf.clip_by_value(actions, -1 + SMALL_NUMBER, 1 - SMALL_NUMBER) actions *= self._action_limit return actions, logpacs
def prepare_t_student_mixture(self): cluster_weights_unnorm = self.get_cluster_unnormalized_weights() cluster_weights_norm = cluster_weights_unnorm / np.sum( cluster_weights_unnorm) cat_distr = tpd.Categorical( probs=tf.cast(cluster_weights_norm, dtype=tf.float32)) t_student_params = [ self.prepare_t_student_params(ind) for ind in range(self.clusters_num) ] dofs = tf.constant( [t_student_params[ind]['df'] for ind in range(self.clusters_num)], dtype=tf.float32) means = tf.stack([ t_student_params[ind]['mean'] for ind in range(self.clusters_num) ]) cov_chols = tf.stack([ t_student_params[ind]['cov_chol'] for ind in range(self.clusters_num) ]) t_student_distr = tpd.MultivariateStudentTLinearOperator( df=dofs, loc=means, scale=tf.linalg.LinearOperatorLowerTriangular(cov_chols)) t_student_mixture = tpd.MixtureSameFamily( mixture_distribution=cat_distr, components_distribution=t_student_distr) return t_student_mixture
def output_function(self, state): params = dense_layer(state.h3, self.output_units, scope='gmm', reuse=tf.compat.v1.AUTO_REUSE) pis, mus, sigmas, rhos, es = self._parse_parameters(params) mu1, mu2 = tf.split(mus, 2, axis=1) mus = tf.stack([mu1, mu2], axis=2) sigma1, sigma2 = tf.split(sigmas, 2, axis=1) covar_matrix = [ tf.square(sigma1), rhos * sigma1 * sigma2, rhos * sigma1 * sigma2, tf.square(sigma2) ] covar_matrix = tf.stack(covar_matrix, axis=2) covar_matrix = tf.reshape( covar_matrix, (self.batch_size, self.num_output_mixture_components, 2, 2)) mvn = tfd.MultivariateNormalFullCovariance( loc=mus, covariance_matrix=covar_matrix) b = tfd.Bernoulli(probs=es) c = tfd.Categorical(probs=pis) sampled_e = b.sample() sampled_coords = mvn.sample() sampled_idx = c.sample() idx = tf.stack([tf.range(self.batch_size), sampled_idx], axis=1) coords = tf.gather_nd(sampled_coords, idx) return tf.concat([coords, tf.cast(sampled_e, tf.float32)], axis=1)
def _train_auxiliary(self, obs, returns, old_pi_logits): if self.policy.is_discrete: old_pi = tfd.Categorical(old_pi_logits) else: mean, logstd = tf.split(old_pi_logits, 2, axis=-1) old_pi = tfd.MultivariateNormalDiag(mean, tf.exp(logstd)) with tf.GradientTape() as tape: pi, aux_value_preds = self.policy.auxiliary_heads(obs) bc_loss = tf.reduce_mean(old_pi.kl_divergence(pi)) aux_value_loss = 0.5 * tf.reduce_mean((returns - aux_value_preds) ** 2) joint_loss = aux_value_loss + self.bc_coef * bc_loss grads = tape.gradient(joint_loss, self.policy.auxiliary_weights) if self.max_grad_norm is not None: grads, _ = tf.clip_by_global_norm(grads, self.max_grad_norm) self.aux_optimizer.apply_gradients(zip(grads, self.policy.auxiliary_weights)) with tf.GradientTape() as tape: value_preds = self.policy.value(obs) value_loss = 0.5 * tf.reduce_mean((returns - value_preds) ** 2) grads = tape.gradient(value_loss, self.policy.value_weights) if self.max_grad_norm is not None: grads, _ = tf.clip_by_global_norm(grads, self.max_grad_norm) if self.value_optimizer is not None: self.value_optimizer.apply_gradients(zip(grads, self.policy.value_weights)) else: self.policy_optimizer.apply_gradients(zip(grads, self.policy.value_weights))
def get_resample_idx(self, log_W, sample_size=()): """ Get resample index a_t^k ~ Categorical(w_t^1, ..., w_t^K) with logits = log_W last axis Input: log_W.shape = (K, batch_size_0, ..., batch_size_last) """ nb_classes = log_W.shape.as_list()[0] batch_shape = log_W.shape.as_list()[1:] perm = list(range(1, len(batch_shape) + 1)) + [0] log_W = tf.transpose(log_W, perm=perm) categorical = tfd.Categorical(logits=log_W, validate_args=True, name="Categorical") # sample multiple times to remove idx out of range if sample_size == (): idx_shape = batch_shape else: assert isinstance(sample_size, int), "sample_size should be int, {} is given".format(sample_size) idx_shape = [sample_size] + batch_shape idx = tf.ones(idx_shape, dtype=tf.int32) * nb_classes for _ in range(1): fixup_idx = categorical.sample(sample_size) idx = tf.where(idx >= nb_classes, fixup_idx, idx) # if still got idx out of range, replace them with idx from uniform distribution final_fixup = tf.random.uniform(idx_shape, maxval=nb_classes, dtype=tf.int32) idx = tf.where(idx >= nb_classes, final_fixup, idx) batch_idx = np.meshgrid(*[range(i) for i in idx_shape], indexing='ij') if sample_size != (): batch_idx = batch_idx[1:] resample_idx = tf.stack([idx] + batch_idx, axis=-1) return resample_idx
def mdn_head(h, FLAGS): with tf.variable_scope('mdn'): locs = tf.reshape(tf.layers.dense(h, 2 * FLAGS['k'], activation=None), [-1, FLAGS['k'], 2]) scales = tf.reshape( tf.layers.dense(h, 2 * FLAGS['k'], activation=tf.exp), [-1, FLAGS['k'], 2]) logits = tf.layers.dense(h, FLAGS['k'], activation=None) cat = tfd.Categorical(logits=logits) components = [] eval_components = [] for loc, scale in zip(tf.unstack(tf.transpose(locs, [1, 0, 2])), tf.unstack(tf.transpose(scales, [1, 0, 2]))): # TODO: does this need to be a more complex distribution? normal = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale) components.append(normal) eval_normal = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale) eval_components.append(eval_normal) mixture = tfd.Mixture(cat=cat, components=components) eval_mixture = tfd.Mixture(cat=cat, components=eval_components) return { 'locs': locs, 'scales': scales, 'logits': logits, 'mixture': mixture, 'eval_mixture': eval_mixture }
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 learn(self): actions = np.array(self.action_memory) # Get actions taken G = self.calc_advantages() # Calc advantages of each action with GradientTape() as tape: # Calculate gradients for all weights loss = 0 for idx, (g, state) in enumerate(zip(G, self.state_memory)): state = convert_to_tensor( [state], dtype='float32') # Makes training faster probs = self.policy(state, training=False)[ 0] # Get probability predictions based on each step state action_probs = distributions.Categorical( probs=probs) # Create distribution based on probabilities log_prob = action_probs.log_prob( actions[idx] ) # Find logistic probability of action being taken loss += -g * squeeze( log_prob ) # Update loss based on reward and model's certainty # Calculate gradients from observing tape, then apply them to the policy weights gradient = tape.gradient(loss, self.policy.trainable_variables) self.policy.optimizer.apply_gradients( zip(gradient, self.policy.trainable_variables)) self.clear_steps() # Clear step memory
def _exploration(self, action, training): if training: amount = self._c.expl_amount if self._c.expl_decay: amount *= 0.5**(tf.cast(self._step, tf.float32) / self._c.expl_decay) if self._c.expl_min: amount = tf.maximum(self._c.expl_min, amount) self._metrics['expl_amount'].update_state(amount) elif self._c.eval_noise: amount = self._c.eval_noise else: return action if self._c.expl == 'additive_gaussian': return tf.clip_by_value(tfd.Normal(action, amount).sample(), -1, 1) if self._c.expl == 'completely_random': return tf.random.uniform(action.shape, -1, 1) if self._c.expl == 'epsilon_greedy': indices = tfd.Categorical(0 * action).sample() rnd = tf.random.uniform(action.shape[:1], 0, 1) < amount rnd = tf.expand_dims(rnd, axis=-1) one_hot = tf.one_hot(indices, action.shape[-1], dtype=self._float) return tf.where(tf.repeat(rnd, 4, axis=-1), one_hot, action) # return tf.where( # tf.random.uniform(action.shape[:1], 0, 1) < amount, # tf.one_hot(indices, action.shape[-1], dtype=self._float), # action) raise NotImplementedError(self._c.expl)
def sample_is(self, x, n=1): mixture_distribution, mixture_components = \ self._gate.conditional_mixture_distribution(x),\ self._experts.conditional_components_distribution(x) y = mixture_components.sample(n) # npdt = y.dtype.as_numpy_dtype is_logits = self._is_function(mixture_distribution.logits) is_mixture_distribution = ds.Categorical(logits=is_logits) idx = is_mixture_distribution.sample(n) # TODO check if we should not renormalize mixture.logits - tf.stop_... weights = tf.batch_gather( mixture_distribution.logits - tf.stop_gradient(is_logits), tf.transpose(idx)) # TODO check axis # weights = tf.batch_gather( # log_normalize(mixture_distribution.logits - tf.stop_gradient(is_logits), axis=1), # tf.transpose(idx)) if n == 1: return tf.batch_gather(y, idx[:, :, None])[0, :, 0], tf.transpose(weights)[0] else: return tf.batch_gather(y, idx[:, :, None])[:, :, 0], tf.transpose(weights)
def get_distribution(self, x, **kwargs): """Build the mixture distribution implied by the set of oracles that are trained in this module Args: x: tf.Tensor a batch of training inputs shaped like [batch_size, channels] Returns: distribution: tfpd.Distribution the mixture of gaussian distributions implied by the oracles """ # get the distribution parameters for all models params = defaultdict(list) for fm in self.forward_models: for key, val in fm.get_params(x, **kwargs).items(): params[key].append(val) # stack the parameters in a new component axis for key, val in params.items(): params[key] = tf.stack(val, axis=-1) # build the mixture distribution using the family of component one weights = tf.fill([self.bootstraps], 1 / self.bootstraps) return tfpd.MixtureSameFamily( tfpd.Categorical(probs=weights), self.forward_models[0].distribution(**params))
def _exploration(self, action, training): if training: amount = self._c.expl_amount if self._c.expl_decay: # 0.3 is True amount *= 0.5**(tf.cast(self._step, tf.float32) / self._c.expl_decay) if self._c.expl_min: # 0.0 is False amount = tf.maximum(self._c.expl_min, amount) self._metrics["expl_amount"].update_state(amount) elif self._c.eval_noise: amount = self._c.eval_noise else: return action if self._c.expl == "additive_gaussian": return tf.clip_by_value(tfd.Normal(action, amount).sample(), -1, 1) if self._c.expl == "completely_random": return tf.random.uniform(action.shape, -1, 1) if self._c.expl == "epsilon_greedy": # print( # "0 * action:", 0 * action # ) # 0 * action: Tensor("mul:0", shape=(1, 4), dtype=float16) indices = tfd.Categorical(0 * action).sample() # print("epsilon_greedy indices:", indices) # shape=(1,) return tf.where( tf.random.uniform(action.shape[:1], 0, 1) < amount, # randomly decide do epsilon greedy or not tf.one_hot(indices, action.shape[-1], dtype=tf.float32), action, ) raise NotImplementedError(self._c.expl)
def mse_func(y_true, y_pred): # Reshape inputs in case this is used in a TimeDistribued layer y_pred = tf.reshape(y_pred, [-1, (2 * num_mixes * output_dim) + num_mixes], name='reshape_ypreds') y_true = tf.reshape(y_true, [-1, output_dim], name='reshape_ytrue') out_mu, out_sigma, out_pi = tf.split(y_pred, num_or_size_splits=[ num_mixes * output_dim, num_mixes * output_dim, num_mixes ], axis=1, name='mdn_coef_split') cat = tfd.Categorical(logits=out_pi) component_splits = [output_dim] * num_mixes 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) ] mixture = tfd.Mixture(cat=cat, components=coll) samp = mixture.sample() mse = tf.reduce_mean(tf.square(samp - y_true), axis=-1) # Todo: temperature adjustment for sampling functon. return mse
def __call__(self): """Get the distribution object from the backend""" if get_backend() == "pytorch": # import torch.distributions as tod raise NotImplementedError else: import tensorflow as tf from tensorflow_probability import distributions as tfd # Convert to tensorflow distributions if probflow distributions if isinstance(self.distributions, BaseDistribution): self.distributions = self.distributions() # Broadcast probs/logits shape = self.distributions.batch_shape args = {"logits": None, "probs": None} if self.logits is not None: args["logits"] = tf.broadcast_to(self["logits"], shape) else: args["probs"] = tf.broadcast_to(self["probs"], shape) # Return TFP distribution object return tfd.MixtureSameFamily( tfd.Categorical(**args), self.distributions )
def sample(self, time, outputs, state, name=None): """ sample tokens for next step, notice the special form of 'state'([decoded_ids, rnn_state]) """ sample_method_sampler = \ tfpd.Categorical(probs=self._lambdas) sample_method_id = sample_method_sampler.sample() truth_feeding = lambda: tf.cond( tf.less(time, tf.shape(self._ground_truth)[1]), lambda: tf.cast(self._ground_truth[:, time], tf.int32), lambda: tf.ones_like(self._ground_truth[:, 0], dtype=tf.int32) * self._vocab.eos_token_id) self_feeding = lambda: SampleEmbeddingHelper.sample( self, time, outputs, state, name) reward_feeding = lambda: self._sample_by_reward(time, state) sample_ids = tf.cond( tf.logical_or(tf.equal(time, 0), tf.equal(sample_method_id, 1)), truth_feeding, lambda: tf.cond( tf.equal(sample_method_id, 2), reward_feeding, self_feeding)) return sample_ids
def get_distributions_from_tensor(t, dimension, num_mixes): y_pred = tf.reshape(t, [-1, (2 * num_mixes * dimension + 1) + num_mixes], name='reshape_ypreds') out_e, out_pi, out_mus, out_stds = tf.split(y_pred, num_or_size_splits=[ 1, num_mixes, num_mixes * dimension, num_mixes * dimension ], name='mdn_coef_split', axis=-1) cat = tfd.Categorical(logits=out_pi) components_splits = [dimension] * num_mixes mus = tf.split(out_mus, num_or_size_splits=components_splits, axis=1) stds = tf.split(out_stds, num_or_size_splits=components_splits, axis=1) components = [ tfd.MultivariateNormalDiag(loc=mu_i, scale_diag=std_i) for mu_i, std_i in zip(mus, stds) ] mix = tfd.Mixture(cat=cat, components=components) stroke = tfd.Bernoulli(logits=out_e) return mix, stroke
def make_mixture_prior(latent_size, mixture_components): """Creates the mixture of Gaussians prior distribution. Args: latent_size: The dimensionality of the latent representation. mixture_components: Number of elements of the mixture. Returns: random_prior: A `tfd.Distribution` instance representing the distribution over encodings in the absence of any evidence. """ if mixture_components == 1: return tfd.MultivariateNormalDiag( loc=tf.zeros([latent_size]), scale_identity_multiplier=1.0) loc = tf.get_variable(name="loc", shape=[mixture_components, latent_size]) raw_scale_diag = tf.get_variable( name="raw_scale_diag", shape=[mixture_components, latent_size]) mixture_logits = tf.get_variable( name="mixture_logits", shape=[mixture_components]) return tfd.MixtureSameFamily( components_distribution=tfd.MultivariateNormalDiag( loc=loc, scale_diag=tf.nn.softplus(raw_scale_diag)), mixture_distribution=tfd.Categorical(logits=mixture_logits), name="prior")
def get_steering(preds): alpha, mu, sigma = slice_parameter_vectors(preds.numpy(), components) # print(alpha) max_prob = np.max(alpha, axis=-1) if max_prob > 0.9995: index = np.argmax(alpha, axis=-1) angle = mu[:, index[0]] else: angle = np.multiply(alpha, mu).sum(axis=-1) # print(angle) gm = tfd.MixtureSameFamily( mixture_distribution=tfd.Categorical(probs=alpha), components_distribution=tfd.Normal(loc=mu, scale=sigma), ) x = np.linspace(-1, 1, int(1e3)) pyx = gm.prob(x) plot = cv2.plot.Plot2d_create( np.array(x).astype(np.float64), np.array(pyx).astype(np.float64)) plot.setPlotBackgroundColor((255, 255, 255)) plot.setInvertOrientation(True) plot.setPlotLineColor(0) plot = plot.render() cv2.imshow("Distribution", plot) return angle[0]