def _build(self, inputs): # create the layers for mean and covariance output_shape = [-1] + self._output_shape logits = tf.reshape(snt.Linear(np.prod(self._output_shape), initializers=self._initializers, regularizers=self._regularizers)(inputs),output_shape) dtype = inputs.dtype if self._dtype is not None: dtype = self._dtype if self._clip_value > 0: probs = tf.nn.sigmoid(logits) probs = tf.clip_by_value(probs, self._clip_value, 1 - self._clip_value) bernoulli = tfd.Bernoulli(probs=probs, dtype=dtype) else: bernoulli = tfd.Bernoulli(logits=logits, dtype=dtype) def reconstruction_node(self): return self.mean() bernoulli.reconstruction_node = types.MethodType(reconstruction_node, bernoulli) def distribution_parameters(self): return [self.mean()] bernoulli.distribution_parameters = types.MethodType(distribution_parameters, bernoulli) def get_probs(self): return self.probs bernoulli.get_probs = types.MethodType(get_probs, bernoulli) return bernoulli
def define_graph(config): network_tpl = tf.make_template('network', network, config=config) inputs = tf.placeholder(tf.float32, [None, config.num_inputs]) targets = tf.placeholder(tf.float32, [None, 1]) num_visible = tf.placeholder(tf.int32, []) batch_size = tf.to_float(tf.shape(inputs)[0]) data_mean, data_noise, data_uncertainty = network_tpl(inputs) ood_inputs = inputs + tf.random_normal( tf.shape(inputs), 0.0, config.noise_std) ood_mean, ood_noise, ood_uncertainty = network_tpl(ood_inputs) losses = [ -tfd.Normal(data_mean, data_noise).log_prob(targets), -tfd.Bernoulli(data_uncertainty).log_prob(0), -tfd.Bernoulli(ood_uncertainty).log_prob(1), ] if config.center_at_target: losses.append(-tfd.Normal(ood_mean, ood_noise).log_prob(targets)) loss = sum(tf.reduce_sum(loss) for loss in losses) / batch_size optimizer = tf.train.AdamOptimizer(config.learning_rate) gradients, variables = zip(*optimizer.compute_gradients( loss, colocate_gradients_with_ops=True)) if config.clip_gradient: gradients, _ = tf.clip_by_global_norm(gradients, config.clip_gradient) optimize = optimizer.apply_gradients(zip(gradients, variables)) data_uncertainty = tf.sigmoid(data_uncertainty) if not config.center_at_target: data_mean = (1 - data_uncertainty) * data_mean + data_uncertainty * 0 data_noise = (1 - data_uncertainty) * data_noise + data_uncertainty * 0.1 return tools.AttrDict(locals())
def call(self, inputs, training=None): logits = self.logits(inputs) q = distributions.Bernoulli( logits=logits, dtype=tf.float32, ) if self.beta > 0.0: kld = q.kl_divergence( distributions.Bernoulli( logits=tf.zeros_like(logits), dtype=tf.float32, )) self.add_loss(tf.reduce_mean(kld)) return q
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 generate_2pl_data(n_sample, n_factor, n_item, alpha, beta, 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) intercept = tf.fill((n_item,), value = tf.constant(alpha, dtype = dtype)) 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) logits = intercept + eta @ tf.transpose(loading) x = tfd.Bernoulli(logits=logits, dtype=dtype).sample() return x
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 __call__(self, inputs): out = self.get('out', tfkl.Dense, np.prod(self._shape))(inputs) out = tf.reshape(out, tf.concat([tf.shape(inputs)[:-1], self._shape], 0)) out = tf.cast(out, tf.float32) if self._dist in ('normal', 'tanh_normal', 'trunc_normal'): std = self.get('std', tfkl.Dense, np.prod(self._shape))(inputs) std = tf.reshape(std, tf.concat([tf.shape(inputs)[:-1], self._shape], 0)) std = tf.cast(std, tf.float32) if self._dist == 'mse': dist = tfd.Normal(out, 1.0) return tfd.Independent(dist, len(self._shape)) if self._dist == 'normal': dist = tfd.Normal(out, std) return tfd.Independent(dist, len(self._shape)) if self._dist == 'binary': dist = tfd.Bernoulli(out) return tfd.Independent(dist, len(self._shape)) if self._dist == 'tanh_normal': mean = 5 * tf.tanh(out / 5) std = tf.nn.softplus(std + self._init_std) + self._min_std dist = tfd.Normal(mean, std) dist = tfd.TransformedDistribution(dist, common.TanhBijector()) dist = tfd.Independent(dist, len(self._shape)) return common.SampleDist(dist) if self._dist == 'trunc_normal': std = 2 * tf.nn.sigmoid((std + self._init_std) / 2) + self._min_std dist = common.TruncNormalDist(tf.tanh(out), std, -1, 1) return tfd.Independent(dist, 1) if self._dist == 'onehot': return common.OneHotDist(out) raise NotImplementedError(self._dist)
def _build(self, inputs): prior = self._prior() enc_out = self._encoder(inputs['s']) loc = self._loc(enc_out) log_scale = self._log_scale(enc_out) scale = tf.nn.softplus( log_scale + softplus_inverse(1.0) ) # idk what this is for. maybe ensuring center around 1.0 q = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale, name='code') # approximate posterior q_sample = q.sample( self.FLAGS['num_vae_samples']) # approximate posterior sample logits = self._decoder_conv(loc) #logits = self._decoder_conv(q_sample) phat = tfd.Independent(tfd.Bernoulli(logits=logits), reinterpreted_batch_ndims=len(IMAGE_SHAPE), name="image") return dict(prior=prior, q=q, q_mean=loc, q_sample=q_sample, phat=phat, phi_s=loc, logits=logits)
def sample(self, time, outputs, state, name=None): """Gets a sample for one step.""" with ops.name_scope(name, "ScheduledOutputTrainingHelperSample", [time, outputs, state]): sampler = tfpd.Bernoulli(probs=self._sampling_probability) return sampler.sample(sample_shape=self.batch_size, seed=self._seed)
def testBernoulliLogProb(self, logits, n): rv = ed.Bernoulli(logits) dist = tfd.Bernoulli(logits) x = rv.distribution.sample(n) rv_log_prob, dist_log_prob = self.evaluate( [rv.distribution.log_prob(x), dist.log_prob(x)]) self.assertAllEqual(rv_log_prob, dist_log_prob)
def decoder(state_sample, observation_dist="gaussian"): """Compute the data distribution of an observation from its state [1].""" check_in( "observation_dist", observation_dist, ("gaussian", "laplace", "bernoulli", "multinomial"), ) timesteps = tf.shape(state_sample)[1] if self.pixel_observations: # original decoder from [1] for deepmind lab envs hidden = tf.layers.dense(state_sample, 1024, None) kwargs = dict(strides=2, activation=tf.nn.relu) hidden = tf.reshape(hidden, [-1, 1, 1, hidden.shape[-1]]) # 1 x 1 hidden = tf.layers.conv2d_transpose(hidden, 128, 5, **kwargs) # 5 x 5 x 128 hidden = tf.layers.conv2d_transpose(hidden, 64, 5, **kwargs) # 13 x 13 x 64 hidden = tf.layers.conv2d_transpose(hidden, 32, 6, **kwargs) # 30 x 30 x 32 mean = 255 * tf.layers.conv2d_transpose( hidden, 3, 6, strides=2, activation=tf.nn.sigmoid) # 64 x 64 x 3 assert mean.shape[1:].as_list() == [64, 64, 3], mean.shape else: # decoder for gridworlds / structured observations hidden = state_sample d = self._hidden_layer_size for _ in range(4): hidden = tf.layers.dense(hidden, d, tf.nn.relu) mean = tf.layers.dense(hidden, np.prod(self.data_shape), None) mean = tf.reshape(mean, [-1, timesteps] + list(self.data_shape)) check_in( "observation_dist", observation_dist, ("gaussian", "laplace", "bernoulli", "multinomial"), ) if observation_dist == "gaussian": dist = tfd.Normal(mean, self._obs_stddev) elif observation_dist == "laplace": dist = tfd.Laplace(mean, self._obs_stddev / np.sqrt(2)) elif observation_dist == "bernoulli": dist = tfd.Bernoulli(probs=mean) else: mean = tf.reshape(mean, [-1, timesteps] + [np.prod(list(self.data_shape))]) dist = tfd.Multinomial(total_count=1, probs=mean) reshape = tfp.bijectors.Reshape( event_shape_out=list(self.data_shape)) dist = reshape(dist) return dist dist = tfd.Independent(dist, len(self.data_shape)) return dist
def __call__(self): """Get the distribution object from the backend""" if get_backend() == 'pytorch': import torch.distributions as tod return tod.bernoulli.Bernoulli(logits=self.logits, probs=self.probs) else: from tensorflow_probability import distributions as tfd return tfd.Bernoulli(logits=self.logits, probs=self.probs)
def _common(cls, node, **kwargs): x = kwargs["tensor_dict"][node.inputs[0]] dtype = node.attrs.get("dtype", x.dtype) dist = tfd.Bernoulli(probs=x, dtype=dtype) if 'seed' in node.attrs: ret = dist.sample(seed=int(node.attrs.get('seed'))) else: ret = dist.sample() return [tf.cast(tf.reshape(ret, x.shape), dtype)]
def __call__(self, features): x = features for index in range(self._layers): x = self.get(f'h{index}', tfkl.Dense, self._units, self._act)(x) x = self.get(f'hout', tfkl.Dense, np.prod(self._shape))(x) x = tf.reshape(x, tf.concat([tf.shape(features)[:-1], self._shape], 0)) if self._dist == 'normal': return tfd.Independent(tfd.Normal(x, 1), len(self._shape)) elif self._dist == 'binary': return tfd.Independent(tfd.Bernoulli(x), len(self._shape)) raise NotImplementedError(self._dist)
def __call__(self, observation, action): x = self._dynamics(tf.concat([observation, action], axis=1)) # TODO (yarden): maybe it's better to feed the reward and terminals s, a instead of x. # The world model predicts the difference between next_observation and observation. return dict(next_observation=tfd.MultivariateNormalDiag( loc=self._next_observation_residual_mu(x) + tf.stop_gradient(observation), scale_diag=self._next_observation_stddev(x)), reward=tfd.Normal(loc=self._reward_mu(x), scale=1.0), terminal=tfd.Bernoulli(logits=self._terminal_logit(x), dtype=tf.float32))
def generate_logistic_reg_data( n_sample, weight, intercept = 0, dtype = tf.float64, seed = None): weight = tf.constant(weight, dtype = dtype) weight = tf.reshape(weight, shape = (-1, 1)) n_feature = weight.shape[0] x = tf.random.normal(shape = (n_sample, n_feature), seed = seed, dtype = dtype) logits = intercept + x @ weight y = tfd.Bernoulli(logits=logits, dtype=dtype).sample() return x, y
def __call__(self, features): kwargs = dict(strides=2, activation=self._act) x = self.get('h1', tfkl.Dense, 8 * self._depth, None)(features) x = tf.reshape(x, [-1, 1, 1, 8 * self._depth]) x = self.get('h2', tfkl.Conv2DTranspose, 4 * self._depth, 5, **kwargs)(x) x = self.get('h3', tfkl.Conv2DTranspose, 2 * self._depth, 5, **kwargs)(x) x = self.get('h4', tfkl.Conv2DTranspose, 1 * self._depth, 6, **kwargs)(x) x = self.get('h5', tfkl.Conv2DTranspose, 1, 6, **kwargs)(x) shape = tf.concat([tf.shape(features)[:-1], self._shape], axis=0) x = tf.reshape(x, shape) return tfd.Independent(tfd.Bernoulli(x), 3) # last 3 dimensions (row, col, chan) define 1 pixel
def sample(self, time, outputs, state, name=None): """Gets a sample for one step.""" with ops.name_scope(name, "ScheduledEmbeddingTrainingHelperSample", [time, outputs, state]): # Return -1s where we did not sample, and sample_ids elsewhere select_sampler = tfpd.Bernoulli(probs=self._sampling_probability, dtype=dtypes.bool) select_sample = select_sampler.sample(sample_shape=self.batch_size, seed=self._scheduling_seed) sample_id_sampler = tfpd.Categorical(logits=outputs) return array_ops.where(select_sample, sample_id_sampler.sample(seed=self._seed), gen_array_ops.fill([self.batch_size], -1))
def call(self, x): x = self._layers(x) if not getattr(self, '_has_cnn', None): rbd = 0 if x.shape[ -1] == 1 else 1 # #reinterpreted batch dimensions x = tf.squeeze(x) if self._dist == 'normal': return tfd.Independent(tfd.Normal(x, 1), rbd) if self._dist == 'binary': return tfd.Independent(tfd.Bernoulli(x), rbd) raise NotImplementedError(self._dist) return x
def decoder(codes): """Builds a distribution over images given codes. Args: codes: A `Tensor` representing the inputs to be decoded, of shape `[..., code_size]`. Returns: decoder_distribution: A multivariate `Bernoulli` distribution. """ logits = decoder_net(codes) return tfd.Independent(tfd.Bernoulli(logits=logits), reinterpreted_batch_ndims=len(output_shape), name="decoder_distribution")
def __call__(self, x): n_sample = len(x) joint_prob = tfd.JointDistributionSequential([ tfd.Independent( tfd.Normal( loc = tf.zeros((n_sample, n_factor), dtype=self.dtype), scale = 1.0), reinterpreted_batch_ndims=1), lambda eta: tfd.Independent( tfd.Bernoulli( logits= self.intercept + eta @ tf.transpose(self.loading), dtype=self.dtype), reinterpreted_batch_ndims=1)]) joint_prob._to_track=self return joint_prob
def __call__(self, features): x = features for index in range(self.layers_num): # 3 x = self.get(f"h{index}", tf.keras.layers.Dense, self._units, self._act)(x) # print("x:",x.shape) # (15, 1250, 400) x = self.get(f"hout", tf.keras.layers.Dense, np.prod(self._shape))(x) # print("x:",x.shape) x = tf.reshape(x, tf.concat([tf.shape(features)[:-1], self._shape], 0)) # print("x:",x.shape) if self._dist == "normal": return tfd.Independent(tfd.Normal(x, 1), len(self._shape)) if self._dist == "binary": return tfd.Independent(tfd.Bernoulli(x), len(self._shape)) raise NotImplementedError(self._dist)
def calculate_log_px_z(x, x_logit, pi): sample_size = x_logit.shape[4] n_required = x_logit.shape[5] # x_broad = tf.repeat(tf.expand_dims(x, 4), axis=4, repeats=pi.shape[1]) # x_broad = tf.reshape(x_broad, shape=(16, 4)) # x_logit = tf.reshape(x_logit, shape=(16, 4)) x_broad = tf.repeat(tf.expand_dims(x, -1), axis=-1, repeats=sample_size) x_broad = tf.repeat(tf.expand_dims(x_broad, -1), axis=-1, repeats=n_required) dist = tfpd.Bernoulli(logits=x_logit) log_px_z = dist.log_prob(x_broad) log_px_z = tf.reduce_sum(log_px_z, axis=(1, 2, 3)) log_px_z = tf.reduce_mean(log_px_z, axis=1) log_px_z = tf.reduce_sum(pi[:, :, 0, 0] * log_px_z, axis=1) log_px_z = tf.reduce_mean(log_px_z, axis=0) return log_px_z
def decode(code, data_shape): """ :param code: number of code units :param data_shape: dimensionality of the input data :return: void """ with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE): logit = tf.reshape( tf.layers.dense( tf.layers.dense( tf.layers.dense(code, 256, tf.nn.relu), 784, tf.nn.relu ), np.prod(data_shape) ), [-1] + data_shape ) return distributions.Independent(distributions.Bernoulli(logit), 2)
def perform_fwd_pass(self): mean, log_var = self.nets.encoder(self.x) stddev = tf.exp(0.5 * log_var) qz_x = tfpd.Normal(loc=mean, scale=stddev) z = qz_x.sample() logits = self.nets.decoder(z) px_z = tfpd.Bernoulli(logits=logits) p_z = tfpd.Normal(loc=tf.zeros_like(z), scale=tf.ones_like(z)) kl = tf.reduce_sum(tfpd.kl_divergence(qz_x, p_z), axis=1) expected_log_likelihood = tf.reduce_sum(px_z.log_prob(self.x), axis=(1, 2, 3)) self.elbo = tf.reduce_mean(expected_log_likelihood - kl, axis=0)
def rollout_in_dream(self, z_init, h_init, video=False): """ Inputs: h_init: (L, B, 1) z_init: (L, B, latent_dim * n_atoms) done_init: (L, B, 1) """ L, B = h_init.shape[:2] horizon = self.config.imagination_horizon z, h = tf.reshape(z_init, [L * B, -1]), tf.reshape(h_init, [L * B, -1]) feats = tf.concat([z, h], axis=-1) #: s_t, a_t, s_t+1 trajectory = {"state": [], "action": [], 'next_state': []} for t in range(horizon): actions = tf.cast(self.policy.sample(feats), dtype=tf.float32) trajectory["state"].append(feats) trajectory["action"].append(actions) h = self.world_model.step_h(z, h, actions) z, _ = self.world_model.rssm.sample_z_prior(h) z = tf.reshape(z, [L * B, -1]) feats = tf.concat([z, h], axis=-1) trajectory["next_state"].append(feats) trajectory = {k: tf.stack(v, axis=0) for k, v in trajectory.items()} #: reward_head(s_t+1) -> r_t #: Distribution.mode()は確立最大値を返すのでNormalの場合は #: trjactory["reward"] == rewards rewards = self.world_model.reward_head(trajectory['next_state']) trajectory["reward"] = rewards disc_logits = self.world_model.discount_head(trajectory['next_state']) trajectory["discount"] = tfd.Independent( tfd.Bernoulli(logits=disc_logits), reinterpreted_batch_ndims=1).mean() return trajectory
def feed_forward( features, data_shape, num_layers=2, activation=tf.nn.relu, mean_activation=None, stop_gradient=False, trainable=True, units=100, std=1.0, low=-1.0, high=1.0, dist='normal', min_std=1e-2, init_std=1.0): hidden = features if stop_gradient: hidden = tf.stop_gradient(hidden) for _ in range(num_layers): hidden = tf.layers.dense(hidden, units, activation, trainable=trainable) mean = tf.layers.dense( hidden, int(np.prod(data_shape)), mean_activation, trainable=trainable) mean = tf.reshape(mean, tools.shape(features)[:-1] + data_shape) if std == 'learned': std = tf.layers.dense( hidden, int(np.prod(data_shape)), None, trainable=trainable) init_std = np.log(np.exp(init_std) - 1) std = tf.nn.softplus(std + init_std) + min_std std = tf.reshape(std, tools.shape(features)[:-1] + data_shape) if dist == 'normal': dist = tfd.Normal(mean, std) dist = tfd.Independent(dist, len(data_shape)) elif dist == 'deterministic': dist = tfd.Deterministic(mean) dist = tfd.Independent(dist, len(data_shape)) elif dist == 'binary': dist = tfd.Bernoulli(mean) dist = tfd.Independent(dist, len(data_shape)) elif dist == 'trunc_normal': # https://www.desmos.com/calculator/rnksmhtgui dist = tfd.TruncatedNormal(mean, std, low, high) dist = tfd.Independent(dist, len(data_shape)) elif dist == 'tanh_normal': # https://www.desmos.com/calculator/794s8kf0es dist = distributions.TanhNormal(mean, std) elif dist == 'tanh_normal_tanh': # https://www.desmos.com/calculator/794s8kf0es mean = 5.0 * tf.tanh(mean / 5.0) dist = distributions.TanhNormal(mean, std) elif dist == 'onehot_score': dist = distributions.OneHot(mean, gradient='score') elif dist == 'onehot_straight': dist = distributions.OneHot(mean, gradient='straight') else: raise NotImplementedError(dist) return dist
def __call__(self, codes): """Build decoder which takes a code and returns a distribution over images. Args: codes: A `float`-like `Tensor` representing the inputs to be decoded. The first dimension (axis 0) indexes batch elements; all other dimensions index event elements. Returns: decoder: A multivariate `Bernoulli` distribution. """ net = self.decoder_net(codes) new_shape = tf.concat([tf.shape(net)[:-1], self.output_shape], axis=0) logits = tf.reshape(net, shape=new_shape) return tfd.Independent(tfd.Bernoulli(logits=logits), reinterpreted_batch_ndims=len( self.output_shape), name="decoder_distribution")
def _compute_log_loss(self, y_true, y_pred, mode): """ Inputs: y_true: (L, B, 1) y_pred: (L, B, 1) mode: "reward" or "discount" """ if mode == "discount": dist = tfd.Independent(tfd.Bernoulli(logits=y_pred), reinterpreted_batch_ndims=1) elif mode == "reward": dist = tfd.Independent(tfd.Normal(loc=y_pred, scale=1.), reinterpreted_batch_ndims=1) log_prob = dist.log_prob(y_true) loss = tf.reduce_mean(log_prob) return loss
def call(self, codes): """Builds a distribution over images given codes. Args: codes: A `Tensor` representing the inputs to be decoded, of shape `[..., code_size]`. Returns: decoder_distribution: A multivariate `Bernoulli` distribution. """ num_samples, batch_size, latent_size, code_size = common_layers.shape_list( codes) codes = tf.reshape(codes, [num_samples * batch_size, latent_size, code_size]) logits = self.decoder_net(codes) logits = tf.reshape(logits, [num_samples, batch_size] + list(self.image_shape)) return tfd.Independent(tfd.Bernoulli(logits=logits), reinterpreted_batch_ndims=len(self.image_shape), name="decoder_distribution")