## Setup codecs
# VAE codec
model = BinaryVAE(hidden_dim=100, latent_dim=40)
model.load_state_dict(torch.load('vae_params'))
rec_net = torch_fun_to_numpy_fun(model.encode)
gen_net = torch_fun_to_numpy_fun(model.decode)
obs_codec = lambda p: cs.Bernoulli(p, bernoulli_precision)
def vae_view(head):
return ag_tuple((np.reshape(head[:latent_size], latent_shape),
np.reshape(head[latent_size:], obs_shape)))
vae_append, vae_pop = cs.repeat(cs.substack(
bb_ans.VAE(gen_net, rec_net, obs_codec, prior_precision, q_precision),
vae_view), num_batches)
## Load mnist images
images = datasets.MNIST(sys.argv[1], train=False, download=True).data.numpy()
images = np.uint64(rng.random_sample(np.shape(images)) < images / 255.)
images = np.split(np.reshape(images, (num_images, -1)), num_batches)
## Encode
# Initialize message with some 'extra' bits
encode_t0 = time.time()
init_message = cs.base_message(obs_size + latent_size)
# Encode the mnist images
message, = vae_append(init_message, images)
message, x = codec.pop(message)
return message, ArraySymbol(x)
return cs.Codec(push, pop)
def vae_view(head):
return ag_tuple(
(np.reshape(head[:latent_size],
latent_shape), np.reshape(head[latent_size:], obs_shape)))
vae_append, vae_pop = cs.Multiset(
cs.substack(
ArrayCodec(
bb_ans.VAE(gen_net, rec_net, obs_codec, prior_precision,
q_precision)), vae_view))
## Load mnist images
images = datasets.MNIST(sys.argv[1], train=False, download=True).data.numpy()
images = np.uint64(rng.random_sample(np.shape(images)) < images / 255.)
images = np.split(np.reshape(images, (num_images, -1)), num_batches)
images = list(map(ArraySymbol, images))
## Encode
# Initialize message with some 'extra' bits
encode_t0 = time.time()
init_message = cs.base_message(obs_size + latent_size)
# Build multiset
multiset = cs.build_multiset(images)