def sample_p(self,
targets_length,
temp,
check_invertibility=False,
targets_mask=None,
**kwargs):
hparams = self._hparams
if targets_mask is None:
targets_mask = ops.sequence_mask(targets_length, hparams)
decoder_self_attention_bias = (
common_attention.attention_bias_ignore_padding(1.0 - targets_mask))
batch_size, targets_max_length = (
common_layers.shape_list(targets_mask)[:2])
prior_shape = [batch_size, targets_max_length, hparams.latent_size]
noise = tf.random.normal(prior_shape, stddev=temp)
p_dist = None
if hparams.prior_type == "standard_normal":
z_p = noise
elif hparams.prior_type == "diagonal_normal":
diag_prior_params = ops.cond_prior("diag_prior", hparams,
tf.zeros(prior_shape),
targets_mask,
hparams.latent_size * 2,
decoder_self_attention_bias,
**kwargs)
p_dist = gops.diagonal_normal(diag_prior_params, "diag_prior")
z_p = p_dist.loc + p_dist.scale * noise
elif hparams.prior_type in ["affine", "additive", "rq"]:
n_levels = len(hparams.depths.split("/"))
divi = max(1, hparams.factor**(n_levels - 1))
flow_prior_shape = [
batch_size, targets_max_length // divi, hparams.latent_size
]
noise = tf.random_normal(flow_prior_shape, stddev=temp)
z_p, _, _, _ = glow.glow("glow",
noise,
targets_mask,
decoder_self_attention_bias,
inverse=True,
init=False,
hparams=self._fparams,
disable_dropout=True,
temp=temp,
**kwargs)
if self.is_evaluating and check_invertibility:
noise_inv, _, _, _ = glow.glow("glow",
z_p,
targets_mask,
decoder_self_attention_bias,
inverse=False,
init=False,
hparams=self._fparams,
disable_dropout=True,
**kwargs)
z_diff = noise - noise_inv
tf.summary.scalar("flow_recon_inverse",
tf.reduce_max(tf.abs(z_diff)))
return z_p, p_dist
def delta_posterior(self, z, targets_mask, decoder_self_attention_bias,
n_gibbs_steps, **kwargs):
hparams = self._hparams
for _ in range(n_gibbs_steps):
_, targets_emb = self.argmax_decode(z, decoder_self_attention_bias,
**kwargs)
q_params = ops.posterior("posterior", hparams, targets_emb,
targets_mask, decoder_self_attention_bias,
**kwargs)
q_dist = gops.diagonal_normal(q_params, "posterior")
z = q_dist.loc # [B, L, C]
return z
def compute_prior_log_prob(
self, z_q, targets_mask, decoder_self_attention_bias,
check_invertibility=False, **kwargs):
hparams = self._hparams
batch_size, targets_max_length = (
common_layers.shape_list(targets_mask)[:2])
prior_shape = [batch_size, targets_max_length, hparams.latent_size]
log_abs_det = tf.zeros([batch_size])
if hparams.prior_type == "standard_normal":
log_p_z_base = gops.standard_normal_density(z_q, targets_mask)
elif hparams.prior_type == "diagonal_normal":
diag_prior_params = ops.cond_prior(
"diag_prior", hparams, tf.zeros(prior_shape), targets_mask,
hparams.latent_size*2, decoder_self_attention_bias, **kwargs)
p_dist = gops.diagonal_normal(diag_prior_params, "diag_prior")
log_p_z_base = p_dist.log_prob(z_q) # [B, L, C]
log_p_z_base = gops.reduce_sum_over_lc(log_p_z_base, targets_mask) # [B]
elif hparams.prior_type in ["affine", "additive", "rq"]:
if self.is_evaluating:
disable_dropout = True
init = False
elif self.is_training:
disable_dropout = False
init = tf.equal(hparams.kl_startup_steps,
tf.cast(tf.train.get_global_step(), tf.int32))
else:
raise ValueError("compute_prior shouldn't be used in decoding.")
z_inv, log_abs_det, log_p_z_base, zs = glow.glow(
"glow", z_q, targets_mask, decoder_self_attention_bias,
inverse=False, init=init, hparams=self._fparams,
disable_dropout=disable_dropout, **kwargs)
if self.is_evaluating and check_invertibility:
z_inv_inv, _, _, _ = glow.glow(
"glow", z_inv, targets_mask, decoder_self_attention_bias,
inverse=True, split_zs=zs, init=False, hparams=self._fparams,
disable_dropout=True, **kwargs)
z_diff = z_q - z_inv_inv
tf.summary.scalar("flow_recon_forward", tf.reduce_max(tf.abs(z_diff)))
return log_p_z_base, log_abs_det
def sample_q(
self, targets, targets_mask, decoder_self_attention_bias, n_samples,
temp, **kwargs):
hparams = self._hparams
batch_size, targets_max_length = common_layers.shape_list(targets_mask)[:2]
q_params = ops.posterior("posterior", hparams, targets, targets_mask,
decoder_self_attention_bias, **kwargs)
q_dist = gops.diagonal_normal(q_params, "posterior")
loc, scale = q_dist.loc, q_dist.scale
z_shape = [batch_size, targets_max_length, hparams.latent_size]
iw_z_shape = [n_samples*batch_size, targets_max_length, hparams.latent_size]
if n_samples == 1:
noise = tf.random_normal(z_shape, stddev=temp)
z_q = loc + scale * noise
log_q_z = q_dist.log_prob(z_q) # [B, L, C]
else:
noise = tf.random_normal([n_samples] + z_shape, stddev=temp)
z_q = loc[tf.newaxis, ...] + scale[tf.newaxis, ...] * noise
log_q_z = q_dist.log_prob(z_q) # [K, B, L, C]
z_q = tf.reshape(z_q, iw_z_shape)
log_q_z = tf.reshape(log_q_z, iw_z_shape)
return z_q, log_q_z, q_dist