def build_model(x, init=False, sampling_mode=False):
global nvp_layers
global nvp_dense_layers
with arg_scope([nn_extra_nvp.conv2d_wn, nn_extra_nvp.dense_wn], init=init):
if len(nvp_layers) == 0:
build_nvp_model()
if len(nvp_dense_layers) == 0:
build_nvp_dense_model()
global student_layer
if student_layer is None:
student_layer = nn_extra_student.StudentRecurrentLayer(
shape=(ndim, ), corr_init=corr_init, nu_init=nu_init)
x_shape = nn_extra_nvp.int_shape(x)
x_bs = tf.reshape(
x, (x_shape[0] * x_shape[1], x_shape[2], x_shape[3], x_shape[4]))
x_bs_shape = nn_extra_nvp.int_shape(x_bs)
log_det_jac = tf.zeros(x_bs_shape[0])
logit_layer = nn_extra_nvp.LogitLayer()
scale_layer = nn_extra_nvp.ScaleLayer()
y, log_det_jac = scale_layer.forward_and_jacobian(
x_bs, None, log_det_jac)
y, log_det_jac = logit_layer.forward_and_jacobian(y, None, log_det_jac)
# construct forward pass
z = None
for layer in nvp_layers:
y, z, log_det_jac = layer.forward_and_jacobian(y, z, log_det_jac)
z = tf.concat([z, y], 3)
for layer in nvp_dense_layers:
z, _, log_det_jac = layer.forward_and_jacobian(
z, None, log_det_jac)
z_shape = nn_extra_nvp.int_shape(z)
z_vec = tf.reshape(z, (x_shape[0], x_shape[1], -1))
log_det_jac = tf.reshape(log_det_jac, (x_shape[0], x_shape[1]))
log_probs = []
z_samples = []
latent_log_probs = []
latent_log_probs_prior = []
with tf.variable_scope("one_step") as scope:
student_layer.reset()
for i in range(seq_len):
if sampling_mode:
z_sample = student_layer.sample(nr_samples=n_samples)
z_samples.append(z_sample)
latent_log_prob = student_layer.get_log_likelihood(
z_sample[:, 0, :])
latent_log_probs.append(latent_log_prob)
else:
latent_log_prob = student_layer.get_log_likelihood(
z_vec[:, i, :])
latent_log_probs.append(latent_log_prob)
log_prob = latent_log_prob + log_det_jac[:, i]
log_probs.append(log_prob)
latent_log_prob_prior = student_layer.get_log_likelihood_under_prior(
z_vec[:, i, :])
latent_log_probs_prior.append(latent_log_prob_prior)
student_layer.update_distribution(z_vec[:, i, :])
scope.reuse_variables()
if sampling_mode:
# one more sample after seeing the last element in the sequence
z_sample = student_layer.sample(nr_samples=n_samples)
z_samples.append(z_sample)
z_samples = tf.concat(z_samples, 1)
latent_log_prob = student_layer.get_log_likelihood(z_sample[:,
0, :])
latent_log_probs.append(latent_log_prob)
z_samples_shape = nn_extra_nvp.int_shape(z_samples)
z_samples = tf.reshape(
z_samples,
(z_samples_shape[0] * z_samples_shape[1], z_shape[1],
z_shape[2], z_shape[3])) # (n_samples*seq_len, z_img_shape)
log_det_jac = tf.zeros(z_samples_shape[0] * z_samples_shape[1])
for layer in reversed(nvp_dense_layers):
z_samples, _, log_det_jac = layer.backward(
z_samples, None, log_det_jac)
x_samples = None
for layer in reversed(nvp_layers):
x_samples, z_samples, log_det_jac = layer.backward(
x_samples, z_samples, log_det_jac)
x_samples, log_det_jac = logit_layer.backward(
x_samples, None, log_det_jac)
x_samples, log_det_jac = scale_layer.backward(
x_samples, None, log_det_jac)
x_samples = tf.reshape(x_samples,
(z_samples_shape[0], z_samples_shape[1],
x_shape[2], x_shape[3], x_shape[4]))
log_det_jac = tf.reshape(log_det_jac,
(z_samples_shape[0], z_samples_shape[1]))
latent_log_probs = tf.stack(latent_log_probs, axis=1)
for i in range(seq_len + 1):
log_prob = latent_log_probs[:, i] - log_det_jac[:, i]
log_probs.append(log_prob)
log_probs = tf.stack(log_probs, axis=1)
return x_samples, log_probs
log_probs = tf.stack(log_probs, axis=1)
latent_log_probs = tf.stack(latent_log_probs, axis=1)
latent_log_probs_prior = tf.stack(latent_log_probs_prior, axis=1)
return log_probs, latent_log_probs, latent_log_probs_prior, z_vec
K = create_covariance_matrix(seq_len, p, x_cov, x_var)
xs = []
for i in range(batch_size):
phi = np.tile(x_mu, (seq_len, 1)).flatten()
x1 = rng.multivariate_normal(phi, K)
x1 = x1.reshape((seq_len, p))
x1 = np.float32(x1)
xs.append(x1[None, :, :])
x = np.concatenate(xs, axis=0)
print('shape x', x.shape)
x_var_tf = tf.placeholder(tf.float32, shape=(batch_size, seq_len, p))
l_rnn = nn_extra_student.StudentRecurrentLayer(shape=(p, ),
nu_init=g_nu,
mu_init=g_mu,
var_init=g_var,
corr_init=g_corr)
l_rnn2 = nn_extra_gauss.GaussianRecurrentLayer(shape=(p, ),
mu_init=g_mu,
var_init=g_var,
corr_init=g_corr)
probs = []
probs_gauss = []
with tf.variable_scope("one_step") as scope:
l_rnn.reset()
l_rnn2.reset()
for i in range(seq_len):
prob_i = l_rnn.get_log_likelihood(x_var_tf[:, i, :])
probs.append(prob_i)