def log_pdf_elbo(self, inputs, v=None): """Construct the jmvae kl objective function. Args: inputs: List of input observations, each a `Tensor` of size `[B, ...]`. v: Placeholder. Returns: jmave_kl_objective: Tensor of size [1], returns the objective value for the evidence lower bound augmented with KL divergence that jmvae-kl computes. """ log_ps, kl_ps = self._log_pdf_elbo_components(inputs, v) log_ps_no_joint = [ p for i, p in enumerate(log_ps) if i != self.joint_distribution_index ] kl_ps_no_joint = [ kl for i, kl in enumerate(kl_ps) if i != self.joint_distribution_index ] # Build evidence lower bound. elbo = tf.reduce_mean( tf.reduce_sum(log_ps_no_joint, axis=0) - kl_ps[self.joint_distribution_index]) kl = tf.reduce_mean(tf.reduce_sum(kl_ps_no_joint, axis=0)) jmvae_kl_objective = elbo - self.jmvae_alpha * kl if self._add_summary: utils.scalar_summary_with_scope('elbo', elbo, self._mode) utils.scalar_summary_with_scope('kl', kl, self._mode) return jmvae_kl_objective
def build_nelbo_loss(self, inputs, mask=None, v=None): """Compute loss for the VAE, supports semi-supervised training. Compute the loss for a minibatch of training samples. Also provides the option to give a mask to specify which terms in the loss to use for each data point, allowing, say semi-supervised learning through the mask. Args: inputs: List of input observations, `Tensor` of size `[B, ...]``. mask: Stacked tensor of [B, 2] `Tensors`, each specifying the weight for a term in the bi-vcca objective. The values in the mask are scaled to lie between 0 to 1 before using. v: `Tensor`, (optional) when conditioning generator on some other modality `v`. Returns: loss: [1] `Tensor` """ bivcca = self.log_pdf_elbos(inputs) if mask is None: batch_size = tf.shape(inputs[0], out_type=tf.int32)[0] mask = tf.ones((batch_size, 2)) mask *= tf.constant([self._mu_tradeoff, 1 - self._mu_tradeoff], dtype=tf.float32) # Normalize each row of the mask to sum to 1 after weighing # the losses. mask /= tf.reduce_sum(mask, axis=1, keep_dims=True) loss_tensor = -1 * tf.reduce_mean(tf.reduce_sum(bivcca * mask, axis=1)) if self._add_summary: utils.scalar_summary_with_scope('NELBO', loss_tensor, self._mode) return loss_tensor
def log_pdf_elbos(self, inputs, v=None): """Compute the evidence lower bound (ELBO). Args: inputs: list of `Tensors`, each an observed modality. v: `Tensor`, (optional) when conditioning generator on some other modality `v`. Returns: multimodal_pdf_elbo: [B, n] `Tensor`, each column has elbo for a given modality, and n = len(inputs). """ log_likelihoods, kls = self._log_pdf_elbo_components(inputs, v) # Build combined evidence lower bound. bivcca_cols = [ tf.reduce_sum([ l for j, l in enumerate(ls) if j != self._joint_distribution_index ], axis=0) - kl for i, (ls, kl) in enumerate(zip(log_likelihoods, kls)) if i != self._joint_distribution_index ] bivcca = tf.stack(bivcca_cols, axis=1, name='concat_bivcca') # Add summaries. if self._add_summary: for i, _ in enumerate(bivcca_cols): utils.scalar_summary_with_scope('bivcca_%d' % i, tf.reduce_mean(bivcca[:, i]), self._mode) return bivcca
def build_nelbo_loss(self, inputs, mask=None, v=None): """Construct the final loss to train the JMVAE. Args: inputs: List of input observations, `Tensor` of size `[B, ...]``. mask: Placeholder, does not do anything. v: Placeholder, does not do anything. Returns: loss: [1] `Tensor`, loss to train the JMVAE. """ if mask is not None: logging.warn('Masking is not implemented for JMVAE.') elbo = self.log_pdf_elbo(inputs) loss_tensor = -elbo if self._add_summary: utils.scalar_summary_with_scope('NELBO', loss_tensor, self._mode) return loss_tensor
def build_nelbo_loss(self, inputs, mask=None, v=None): """Compute loss for the VAE, supports semi-supervised training. Compute the loss for a minibatch of training samples. Also provides the option to give a mask to specify which terms in the loss to use for each data point, allowing, say semi-supervised learning through the mask. Args: inputs: List of input observations, `Tensor` of size `[B, ...]``. mask: [B, 3] `Tensor`, mask specifying if x, y or both x and y are present. For example, if the training data point just has x, the mask would be [1.0 0.0 0.0], for x and y it would be [1.0 1.0 1.0]. v: `Tensor`, (optional) when conditioning generator on some other modality `v`. Returns: multimodal_df_elbo: [B, 3] `Tensor`, each column has elbo for [x, y, xy] """ x, y, _ = inputs multimodal_pdf_elbo = self.log_pdf_elbos(x, y) if self._stop_gradient_y is 'REPLACE_ME': return self.build_sg_loss(x, y) if mask is None: mask = tf.ones((tf.shape(x, out_type=tf.int32)[0], 3)) # Normalize each row of the mask to sum to 1 after weighing # the losses. mask /= tf.reduce_sum(mask, axis=1, keep_dims=True) loss_tensor = -1 * tf.reduce_mean( tf.reduce_sum(multimodal_pdf_elbo * mask, axis=1)) if self._add_summary: utils.scalar_summary_with_scope('NELBO', loss_tensor, self._mode) return loss_tensor
def log_pdf_elbos(self, x, y, v=None): """Compute the evidence lower bound (ELBO). Args: x: `Tensor`, observed modality 'x'. y: `Tensor`, observed modality 'y'. v: `Tensor`, (optional) when conditioning generator on some other modality `v`. Returns: multimodal_pdf_elbo: [B, 3] `Tensor`, each column has elbo for [x, y, xy] """ ll, kl = self._log_pdf_elbo_components(x, y, v) # Build combined evidence lower bound. if self._predict_both: multimodal_pdf_elbo = tf.stack( (ll.log_p_x_qx + ll.log_p_y_qx - kl.kl_qx_p, ll.log_p_y_qy + ll.log_p_x_qy - kl.kl_qy_p, ll.log_p_x_qxy + ll.log_p_y_qxy - kl.kl_qxy_p), axis=1, name='concat_elbo') else: multimodal_pdf_elbo = tf.stack( (ll.log_p_x_qx - kl.kl_qx_p, ll.log_p_y_qy - kl.kl_qy_p, ll.log_p_x_qxy + ll.log_p_y_qxy - kl.kl_qxy_p), axis=1, name='concat_elbo') # Add summaries. if self._add_summary: utils.scalar_summary_with_scope( 'elbo_x', tf.reduce_mean(multimodal_pdf_elbo[:, 0]), self._mode) utils.scalar_summary_with_scope( 'elbo_y', tf.reduce_mean(multimodal_pdf_elbo[:, 1]), self._mode) utils.scalar_summary_with_scope( 'elbo_xy', tf.reduce_mean(multimodal_pdf_elbo[:, 2]), self._mode) return multimodal_pdf_elbo