def evaluate_loop(dataset, FactorMuE_variational, latent_dims, latent_length,
latent_alphabet_size, alphabet_size, transfer_mats, bt_scale,
b0_scale, u_conc, r_conc, l_conc, z_distr, mc_samples, dtype,
writer):
"""Evaluate heldout log likelihood and perplexity."""
data_size = sum(1 for i in dataset.batch(1))
heldout_likelihood = 0.
heldout_log_perplex = 0.
for x_batch, xlen_batch in dataset.batch(1):
x = x_batch[0]
xlen = xlen_batch[0]
log_likelihood = 0.
for rep in range(mc_samples):
# Sample from variational approximation.
qz, qbt, qb0, qu, qr, ql = FactorMuE_variational(x)
# Forward pass.
with tape() as model_tape:
# Condition on variational sample.
with condition(z=qz, bt=qbt, b0=qb0, u=qu, r=qr, l=ql):
posterior_predictive = FactorMuE(latent_dims,
latent_length,
latent_alphabet_size,
alphabet_size,
xlen,
transfer_mats,
bt_scale,
b0_scale,
u_conc,
r_conc,
l_conc,
z_distr=z_distr,
dtype=dtype)
# Compute likelihood term.
log_likelihood += mue.hmm_log_prob(
posterior_predictive.distribution, x, xlen) / mc_samples
# Compute KL(vi posterior||prior) for local parameters.
kl_local = 0.
for rv_name, variational_rv in [("z", qz)]:
kl_local += tf.reduce_sum(
variational_rv.distribution.kl_divergence(
model_tape[rv_name].distribution))
# Summary
local_elbo = log_likelihood - kl_local
heldout_likelihood += local_elbo
heldout_log_perplex -= local_elbo / tf.cast(xlen * data_size, dtype)
# Compute perplexity.
heldout_perplex = tf.exp(heldout_log_perplex)
# Record.
with writer.as_default():
tf.summary.scalar('perplexity', heldout_perplex, step=1)
return heldout_perplex, heldout_likelihood
def project_latent_to_sequence(zs, RegressMuE_variational, latent_dims,
latent_length, latent_alphabet_size,
alphabet_size, bt_scale, b0_scale, u_conc,
r_conc, l_conc, z_covar=None, x=None, xlen=None,
mc_samples=1, dtype=tf.float32):
"""Project from latent space to sequence space."""
# Make transfer matrices.
transfer_mats = mue.make_transfer(latent_length, dtype=dtype)
if x is not None:
# Get alignment projection matrix.
pmode_oh = single_alignment_mode(
RegressMuE_variational, z_covar, x, xlen, latent_dims,
latent_length, latent_alphabet_size, alphabet_size,
transfer_mats, bt_scale, b0_scale, u_conc, r_conc,
l_conc, mc_samples=mc_samples, dtype=dtype)
else:
# Project onto conserved positions.
pmode_oh = tf.convert_to_tensor(
mue.mg2k(np.arange(latent_length), 0)[:, None]
== np.arange(2*latent_length + 2)[None, :], dtype=dtype)
xlen = latent_length
# Store results.
nus = [tf.zeros([xlen, alphabet_size], dtype=dtype)
for iz, z in enumerate(zs)]
for rep in range(mc_samples):
# Sample from variational approximation.
qbt, qb0, qu, qr, ql = RegressMuE_variational()
for iz, z in enumerate(zs):
# Condition on latent space position.
with condition(bt=qbt, b0=qb0, u=qu, r=qr, l=ql):
posterior_predictive = RegressMuE(
z, latent_dims, latent_length,
latent_alphabet_size, alphabet_size,
xlen, transfer_mats, bt_scale, b0_scale,
u_conc, r_conc, l_conc, dtype=dtype)
# Compute latent-sequence space observation.
latseq = tf.exp(
posterior_predictive.distribution.observation_distribution.logits)
# Project to sequence space.
nus[iz] += tf.matmul(pmode_oh, latseq) / mc_samples
return nus
def evaluate_loop(dataset, RegressMuE_variational,
latent_dims, latent_length,
latent_alphabet_size, alphabet_size,
transfer_mats, bt_scale, b0_scale, u_conc, r_conc, l_conc,
mc_samples, dtype, writer):
"""Evaluate heldout log likelihood and perplexity."""
data_size = sum(1 for i in dataset.batch(1))
heldout_likelihood = 0.
heldout_log_perplex = 0.
for z_batch, x_batch, xlen_batch in dataset.batch(1):
z = z_batch[0]
x = x_batch[0]
xlen = xlen_batch[0]
log_likelihood = 0.
for rep in range(mc_samples):
# Sample from variational approximation.
qbt, qb0, qu, qr, ql = RegressMuE_variational()
# Condition on variational sample.
with condition(bt=qbt, b0=qb0, u=qu, r=qr, l=ql):
posterior_predictive = RegressMuE(
z, latent_dims, latent_length,
latent_alphabet_size, alphabet_size,
xlen, transfer_mats, bt_scale, b0_scale,
u_conc, r_conc, l_conc, dtype=dtype)
# Compute likelihood term.
log_likelihood += mue.hmm_log_prob(
posterior_predictive.distribution, x, xlen) / mc_samples
# Summary
local_elbo = log_likelihood
heldout_likelihood += local_elbo
heldout_log_perplex -= local_elbo/tf.cast(xlen * data_size, dtype)
# Compute perplexity.
heldout_perplex = tf.exp(heldout_log_perplex)
# Record.
with writer.as_default():
tf.summary.scalar('perplexity', heldout_perplex, step=1)
return heldout_perplex, heldout_likelihood
def single_alignment_mode(RegressMuE_variational, z, x, xlen, latent_dims,
latent_length, latent_alphabet_size, alphabet_size,
transfer_mats, bt_scale, b0_scale, u_conc, r_conc,
l_conc, mc_samples=1, z_distr='Normal',
dtype=tf.float32):
"""Align example sequence (eg. for plotting on structure)."""
# Sample from variational approximation.
qbt, qb0, qu, qr, ql = RegressMuE_variational()
pmode_tuples = []
for rep in range(mc_samples):
# Condition on variational sample.
with condition(bt=qbt.distribution.sample(),
b0=qb0.distribution.sample(),
u=qu.distribution.sample(),
r=qr.distribution.sample(), l=ql.distribution.sample()):
posterior_predictive = RegressMuE(
z, latent_dims, latent_length,
latent_alphabet_size, alphabet_size,
xlen, transfer_mats, bt_scale, b0_scale,
u_conc, r_conc, l_conc, dtype=dtype)
# Compute posterior mode.
pmode = posterior_predictive.distribution.posterior_mode(x[:xlen, :])
# Convert to tuple.
pmode_tuples.append(tuple(pmode.numpy()))
# Get mode (MAP).
pmode_MAP = mue.get_most_common_tuple(pmode_tuples)
print('{} unique alignment(s) out of {} Monte Carlo samples'.format(
len(set(pmode_tuples)), mc_samples))
# One hot encode (seq position x hidden state)
pmode_oh = (tf.convert_to_tensor(pmode_MAP, dtype=tf.int32)[:, None]
== tf.range(
posterior_predictive.distribution.observation_distribution.logits.shape[
0], dtype=tf.int32)[None, :])
return tf.cast(pmode_oh, dtype=dtype)
def train_loop(dataset, RegressMuE_variational, trainable_variables,
latent_dims, latent_length, latent_alphabet_size, alphabet_size,
transfer_mats, bt_scale, b0_scale, u_conc, r_conc, l_conc,
max_epochs, shuffle_buffer, batch_size,
optimizer_name, learning_rate, dtype, writer, out_folder):
"""Training loop."""
# Set up the optimizer.
optimizer = getattr(tf.keras.optimizers, optimizer_name)(
learning_rate=learning_rate)
data_size = sum(1 for i in dataset.batch(1))
# Training loop.
step = 0
t0 = datetime.datetime.now()
for epoch in range(max_epochs):
for z_batch, x_batch, xlen_batch in dataset.shuffle(
shuffle_buffer).batch(batch_size):
step += 1
accum_gradients = [tf.zeros(elem.shape, dtype=dtype)
for elem in trainable_variables]
accum_elbo = 0.
ix = -1
for z, x, xlen in zip(z_batch, x_batch, xlen_batch):
ix += 1
# Track gradients.
with tf.GradientTape() as gtape:
# Sample from variational approximation.
qbt, qb0, qu, qr, ql = RegressMuE_variational()
# Forward pass.
with tape() as model_tape:
# Condition on variational sample.
with condition(bt=qbt, b0=qb0, u=qu, r=qr, l=ql):
posterior_predictive = RegressMuE(
z, latent_dims, latent_length,
latent_alphabet_size, alphabet_size,
xlen, transfer_mats, bt_scale, b0_scale,
u_conc, r_conc, l_conc, dtype=dtype)
# Compute likelihood term.
log_likelihood = mue.hmm_log_prob(
posterior_predictive.distribution, x, xlen)
# Compute KL(vi posterior||prior) for global parameters.
kl_global = 0.
if ix == x_batch.shape[0] - 1:
for rv_name, variational_rv in [
("bt", qbt), ("b0", qb0),
("u", qu), ("r", qr), ("l", ql)]:
kl_global += tf.reduce_sum(
variational_rv.distribution.kl_divergence(
model_tape[rv_name].distribution))
# Compute the ELBO term, correcting for subsampling.
elbo = ((data_size/x_batch.shape[0])
* log_likelihood - kl_global)
# Compute gradient.
loss = -elbo
gradients = gtape.gradient(loss, trainable_variables)
# Accumulate elbo and gradients.
accum_elbo += elbo
for gi, grad in enumerate(gradients):
if grad is not None:
accum_gradients[gi] += grad
# Record.
with writer.as_default():
tf.summary.scalar('elbo', accum_elbo, step=step)
# Optimization step.
optimizer.apply_gradients(
zip(accum_gradients, trainable_variables))
print('epoch {} ({})'.format(epoch, datetime.datetime.now() - t0))
return RegressMuE_variational, trainable_variables