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