def posterior_pop(message):
# pop top-down
(post_mean, post_stdd), h_rec = rec_net_top(contexts[-1])
(prior_mean, prior_stdd), h_gen = gen_net_top()
_, pop = codecs.substack(
codecs.DiagGaussian_GaussianBins(post_mean, post_stdd,
prior_mean, prior_stdd,
latent_prec, prior_prec),
z_view)
message, latent = pop(message)
latents = [(latent, (prior_mean, prior_stdd))]
for rec_net, gen_net, context in reversed(
list(zip(rec_nets, gen_nets, contexts[:-1]))):
previous_latent_val = prior_mean + \
codecs.std_gaussian_centres(prior_prec)[latents[-1][0]] * prior_stdd
(post_mean,
post_stdd), h_rec = rec_net(h_rec, previous_latent_val,
context)
(prior_mean,
prior_stdd), h_gen = gen_net(h_gen, previous_latent_val)
_, pop = codecs.substack(
codecs.DiagGaussian_GaussianBins(post_mean, post_stdd,
prior_mean, prior_stdd,
latent_prec, prior_prec),
z_view)
message, latent = pop(message)
latents.append((latent, (prior_mean, prior_stdd)))
return message, (latents[::-1], h_gen)
def posterior(data):
mu1, sig1, h = rec_net1(data)
mu2, sig2 = rec_net2(h)
post_z2_append, post_z2_pop = codecs.substack(
DiagGaussian_StdBins(mu2, sig2, latent_prec, prior_prec), z2_view)
def posterior_append(message, latents):
(z1, z2), theta1 = latents
z2_vals = codecs.std_gaussian_centres(prior_prec)[z2]
post_z1_append, _ = codecs.substack(
post1_codec(z2_vals, mu1, sig1), z1_view)
theta1[..., 0] = mu1
theta1[..., 1] = sig1
message = post_z1_append(message, z1, theta1)
message = post_z2_append(message, z2)
return message
def posterior_pop(message):
message, z2 = post_z2_pop(message)
z2_vals = codecs.std_gaussian_centres(prior_prec)[z2]
# need to return theta1 from the z1 pop
_, post_z1_pop = codecs.substack(post1_codec(z2_vals, mu1, sig1),
z1_view)
message, (z1, theta1) = post_z1_pop(message)
return message, ((z1, z2), theta1)
return posterior_append, posterior_pop
def posterior_append(message, latents):
# first run the model top-down to get the params and latent vals
latents, _ = latents
(post_mean, post_stdd), h_rec = rec_net_top(contexts[-1])
post_params = [(post_mean, post_stdd)]
for rec_net, latent, context in reversed(
list(zip(rec_nets, latents[1:], contexts[:-1]))):
previous_latent, (prior_mean, prior_stdd) = latent
previous_latent_val = prior_mean + \
codecs.std_gaussian_centres(prior_prec)[previous_latent] * prior_stdd
(post_mean,
post_stdd), h_rec = rec_net(h_rec, previous_latent_val,
context)
post_params.append((post_mean, post_stdd))
# now append bottom up
for latent, post_param in zip(latents, reversed(post_params)):
latent, (prior_mean, prior_stdd) = latent
post_mean, post_stdd = post_param
append, _ = codecs.substack(
codecs.DiagGaussian_GaussianBins(post_mean, post_stdd,
prior_mean, prior_stdd,
latent_prec, prior_prec),
z_view)
message = append(message, latent)
return message
def likelihood(latents):
# get the z1 vals to condition on
latents, h = latents
z1_idxs, (prior_mean, prior_stdd) = latents[0]
z1_vals = prior_mean + codecs.std_gaussian_centres(
prior_prec)[z1_idxs] * prior_stdd
return codecs.substack(obs_codec(h, z1_vals), x_view)
def posterior_pop(message):
message, z2 = post_z2_pop(message)
z2_vals = codecs.std_gaussian_centres(prior_prec)[z2]
# need to return theta1 from the z1 pop
_, post_z1_pop = codecs.substack(post1_codec(z2_vals, mu1, sig1),
z1_view)
message, (z1, theta1) = post_z1_pop(message)
return message, ((z1, z2), theta1)
def likelihood(latent):
(z1, _), theta1 = latent
# get z1_vals from the latent
_, _, mu1_prior, sig1_prior = np.moveaxis(theta1, -1, 0)
eps1_vals = codecs.std_gaussian_centres(prior_prec)[z1]
z1_vals = mu1_prior + sig1_prior * eps1_vals
append, pop = codecs.substack(obs_codec(gen_net2_partial(z1_vals)),
x_view)
return append, pop
def posterior_append(message, latents):
(z1, z2), theta1 = latents
z2_vals = codecs.std_gaussian_centres(prior_prec)[z2]
post_z1_append, _ = codecs.substack(
post1_codec(z2_vals, mu1, sig1), z1_view)
theta1[..., 0] = mu1
theta1[..., 1] = sig1
message = post_z1_append(message, z1, theta1)
message = post_z2_append(message, z2)
return message
def VAE(gen_net, rec_net, obs_codec, prior_prec, latent_prec):
"""
This codec uses the BB-ANS algorithm to code data which is distributed
according to a variational auto-encoder (VAE) model. It is assumed that the
VAE uses an isotropic Gaussian prior and diagonal Gaussian for its
posterior.
"""
z_view = lambda head: head[0]
x_view = lambda head: head[1]
prior = substack(Uniform(prior_prec), z_view)
def likelihood(latent_idxs):
z = std_gaussian_centres(prior_prec)[latent_idxs]
return substack(obs_codec(gen_net(z)), x_view)
def posterior(data):
post_mean, post_stdd = rec_net(data)
return substack(DiagGaussian_StdBins(
post_mean, post_stdd, latent_prec, prior_prec), z_view)
return BBANS(prior, likelihood, posterior)
def codec_from_shape(shape):
print("Creating codec for shape " + str(shape))
hps.image_size = (shape[2], shape[3])
z_shape = latent_shape(hps)
z_size = np.prod(z_shape)
graph = tf.Graph()
with graph.as_default():
with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
x = tf.placeholder(tf.float32, shape, 'x')
model = CVAE1(hps, "eval", x)
stepwise_model = LayerwiseCVAE(model)
saver = tf.train.Saver(model.avg_dict)
config = tf.ConfigProto(allow_soft_placement=True,
intra_op_parallelism_threads=4,
inter_op_parallelism_threads=4,
device_count={'GPU': 0})
sess = tf.Session(config=config, graph=graph)
saver.restore(sess, restore_path())
run_all_contexts, run_top_prior, runs_down_prior, run_top_posterior, runs_down_posterior, \
run_reconstruction = stepwise_model.get_model_parts_as_numpy_functions(sess)
# Setup codecs
def vae_view(head):
return ag_tuple(
(np.reshape(head[:z_size],
z_shape), np.reshape(head[z_size:], shape)))
obs_codec = lambda h, z1: codecs.Logistic_UnifBins(
*run_reconstruction(h, z1), obs_precision, bin_prec=8)
return codecs.substack(
ResNetVAE(run_all_contexts, run_top_posterior, runs_down_posterior,
run_top_prior, runs_down_prior, obs_codec,
prior_precision, q_precision), vae_view)
def SamplingWithoutReplacement():
'''
Encodes and pops onto the ANS state using the empirical
distribution of symbols in the multiset.
Before an push, the symbol to be pushd is inserted into the multiset.
After a pop, the popd symbol is removed from the multiset. Therefore,
a pop performs sampling without replacement, while push inverts it.
The context is the multiset, i.e. *context = multiset
'''
def push(ans_state, symbol, multiset):
multiset, (start, freq) = insert_then_forward_lookup(multiset, symbol)
multiset_size = multiset[0]
ans_state = rans_push(ans_state, start, freq, multiset_size)
return ans_state, multiset
def pop(ans_state, multiset):
multiset_size = multiset[0]
cdf_value, pop_ = rans_pop(ans_state, multiset_size)
multiset, (start, freq), symbol = \
reverse_lookup_then_remove(multiset, cdf_value[0])
ans_state = pop_(start, freq)
return ans_state, symbol, multiset
return substack(Codec(push, pop), lambda head: head[:1])
def posterior(data):
post_mean, post_stdd = rec_net(data)
return substack(
DiagGaussian_StdBins(post_mean, post_stdd, latent_prec,
prior_prec), z_view)
def likelihood(latent_idxs):
z = std_gaussian_centres(prior_prec)[latent_idxs]
return substack(obs_codec(gen_net(z)), x_view)
def ResNetVAE(up_pass, rec_net_top, rec_nets, gen_net_top, gen_nets, obs_codec,
prior_prec, latent_prec):
"""
Codec for a ResNetVAE.
Assume that the posterior is bidirectional -
i.e. has a deterministic upper pass but top down sampling.
Further assume that all latent conditionals are factorised Gaussians,
both in the generative network p(z_n|z_{n-1})
and in the inference network q(z_n|x, z_{n-1})
Assume that everything is ordered bottom up
"""
z_view = lambda head: head[0]
x_view = lambda head: head[1]
prior_codec = codecs.substack(codecs.Uniform(prior_prec), z_view)
def prior_append(message, latents):
# append bottom-up
append, _ = prior_codec
latents, _ = latents
for latent in latents:
latent, _ = latent
message = append(message, latent)
return message
def prior_pop(message):
# pop top-down
(prior_mean, prior_stdd), h_gen = gen_net_top()
_, pop = prior_codec
message, latent = pop(message)
latents = [(latent, (prior_mean, prior_stdd))]
for gen_net in reversed(gen_nets):
previous_latent_val = prior_mean + codecs.std_gaussian_centres(
prior_prec)[latent] * prior_stdd
(prior_mean, prior_stdd), h_gen = gen_net(h_gen,
previous_latent_val)
message, latent = pop(message)
latents.append((latent, (prior_mean, prior_stdd)))
return message, (latents[::-1], h_gen)
def posterior(data):
# run deterministic upper-pass
contexts = up_pass(data)
def posterior_append(message, latents):
# first run the model top-down to get the params and latent vals
latents, _ = latents
(post_mean, post_stdd), h_rec = rec_net_top(contexts[-1])
post_params = [(post_mean, post_stdd)]
for rec_net, latent, context in reversed(
list(zip(rec_nets, latents[1:], contexts[:-1]))):
previous_latent, (prior_mean, prior_stdd) = latent
previous_latent_val = prior_mean + \
codecs.std_gaussian_centres(prior_prec)[previous_latent] * prior_stdd
(post_mean,
post_stdd), h_rec = rec_net(h_rec, previous_latent_val,
context)
post_params.append((post_mean, post_stdd))
# now append bottom up
for latent, post_param in zip(latents, reversed(post_params)):
latent, (prior_mean, prior_stdd) = latent
post_mean, post_stdd = post_param
append, _ = codecs.substack(
codecs.DiagGaussian_GaussianBins(post_mean, post_stdd,
prior_mean, prior_stdd,
latent_prec, prior_prec),
z_view)
message = append(message, latent)
return message
def posterior_pop(message):
# pop top-down
(post_mean, post_stdd), h_rec = rec_net_top(contexts[-1])
(prior_mean, prior_stdd), h_gen = gen_net_top()
_, pop = codecs.substack(
codecs.DiagGaussian_GaussianBins(post_mean, post_stdd,
prior_mean, prior_stdd,
latent_prec, prior_prec),
z_view)
message, latent = pop(message)
latents = [(latent, (prior_mean, prior_stdd))]
for rec_net, gen_net, context in reversed(
list(zip(rec_nets, gen_nets, contexts[:-1]))):
previous_latent_val = prior_mean + \
codecs.std_gaussian_centres(prior_prec)[latents[-1][0]] * prior_stdd
(post_mean,
post_stdd), h_rec = rec_net(h_rec, previous_latent_val,
context)
(prior_mean,
prior_stdd), h_gen = gen_net(h_gen, previous_latent_val)
_, pop = codecs.substack(
codecs.DiagGaussian_GaussianBins(post_mean, post_stdd,
prior_mean, prior_stdd,
latent_prec, prior_prec),
z_view)
message, latent = pop(message)
latents.append((latent, (prior_mean, prior_stdd)))
return message, (latents[::-1], h_gen)
return posterior_append, posterior_pop
def likelihood(latents):
# get the z1 vals to condition on
latents, h = latents
z1_idxs, (prior_mean, prior_stdd) = latents[0]
z1_vals = prior_mean + codecs.std_gaussian_centres(
prior_prec)[z1_idxs] * prior_stdd
return codecs.substack(obs_codec(h, z1_vals), x_view)
return BBANS((prior_append, prior_pop), likelihood, posterior)
obs_elem_codec = lambda p, idx: codecs.Categorical(p, obs_precision)
obs_codec = lambda theta: codecs.AutoRegressive(
obs_elem_param_fn(theta), np.shape(images[0]),
np.shape(images[0]) + (256, ), obs_elem_idxs, obs_elem_codec)
# Setup codecs
def vae_view(head):
return ag_tuple(
(np.reshape(head[:latent1_size], latent1_shape),
np.reshape(head[latent1_size:latent1_size + latent2_size],
latent2_shape),
np.reshape(head[latent1_size + latent2_size:], (batch_size, ))))
vae_append, vae_pop = codecs.repeat(
codecs.substack(
TwoLayerVAE(gen_net2_partial, rec_net1, rec_net2, post1_codec,
obs_codec, prior_precision, q_precision, get_theta1),
vae_view), num_batches)
other_bits_count = 1000000
init_message = codecs.random_stack(
other_bits_count, batch_size + latent1_size + latent2_size, rng)
init_len = 32 * other_bits_count
encode_t0 = time.time()
message = vae_append(init_message, images.astype('uint64'))
flat_message = codecs.flatten(message)
encode_t = time.time() - encode_t0
print("All encoded in {:.2f}s".format(encode_t))
def TwoLayerVAE(gen_net2_partial, rec_net1, rec_net2, post1_codec, obs_codec,
prior_prec, latent_prec, get_theta):
"""
rec_net1 outputs params for q(z1|x)
rec_net2 outputs params for q(z2|x)
post1_codec is to code z1 by q(z1|z2,x)
obs_codec is to code x by p(x|z1)"""
z1_view = lambda head: head[0]
z2_view = lambda head: head[1]
x_view = lambda head: head[2]
prior_z1_append, prior_z1_pop = codecs.substack(Uniform(prior_prec),
z1_view)
prior_z2_append, prior_z2_pop = codecs.substack(Uniform(prior_prec),
z2_view)
def prior_append(message, latent):
(z1, z2), theta1 = latent
message = prior_z1_append(message, z1)
message = prior_z2_append(message, z2)
return message
def prior_pop(message):
message, z2 = prior_z2_pop(message)
message, z1 = prior_z1_pop(message)
# compute theta1
eps1_vals = codecs.std_gaussian_centres(prior_prec)[z1]
z2_vals = codecs.std_gaussian_centres(prior_prec)[z2]
theta1 = get_theta(eps1_vals, z2_vals)
return message, ((z1, z2), theta1)
def likelihood(latent):
(z1, _), theta1 = latent
# get z1_vals from the latent
_, _, mu1_prior, sig1_prior = np.moveaxis(theta1, -1, 0)
eps1_vals = codecs.std_gaussian_centres(prior_prec)[z1]
z1_vals = mu1_prior + sig1_prior * eps1_vals
append, pop = codecs.substack(obs_codec(gen_net2_partial(z1_vals)),
x_view)
return append, pop
def posterior(data):
mu1, sig1, h = rec_net1(data)
mu2, sig2 = rec_net2(h)
post_z2_append, post_z2_pop = codecs.substack(
DiagGaussian_StdBins(mu2, sig2, latent_prec, prior_prec), z2_view)
def posterior_append(message, latents):
(z1, z2), theta1 = latents
z2_vals = codecs.std_gaussian_centres(prior_prec)[z2]
post_z1_append, _ = codecs.substack(
post1_codec(z2_vals, mu1, sig1), z1_view)
theta1[..., 0] = mu1
theta1[..., 1] = sig1
message = post_z1_append(message, z1, theta1)
message = post_z2_append(message, z2)
return message
def posterior_pop(message):
message, z2 = post_z2_pop(message)
z2_vals = codecs.std_gaussian_centres(prior_prec)[z2]
# need to return theta1 from the z1 pop
_, post_z1_pop = codecs.substack(post1_codec(z2_vals, mu1, sig1),
z1_view)
message, (z1, theta1) = post_z1_pop(message)
return message, ((z1, z2), theta1)
return posterior_append, posterior_pop
return BBANS((prior_append, prior_pop), likelihood, posterior)
