def pop(ans: ANSCoder):
symbols = []
for i in range(0, n):
cf, pop = ans.pop(precision)
start = cf
freq = 1
pop(ans, start, freq)
symbols.append(cf)
return symbols
def pop(ans: ANSCoder):
symbols = []
for ppf, cdf in zip(ppdfs, cdfs):
bits, pop_symbol = ans.pop(precision)
index = ppf(bits)
start = cdf(index)
frequency = cdf(index + 1) - start
pop_symbol(ans, start, frequency)
symbols.append(index)
return np.array(symbols)
def beta_binomial_experiment(number_images, result_path, latent_dim,
hidden_dim, model, path):
"""
Run a beta binomial experiment with the given model
"""
seed = 0
rng = RandomState(seed)
image_count = number_images
prior_precision = 8
beta_binomial_precision = 14
q_precision = 14
random_count = 50
number_buckets = 255 #(trials in the beta-binomial distribution, parameter n. It has to be less than 1 << beta_binomial_precision)
ans = ANSCoder()
# obtain numpy compatible functions
generative_model = torch_to_numpy_function(model.decode)
recognition_model = torch_to_numpy_function(model.encode)
latent_shape = (1, latent_dim)
append, pop = build_beta_binomial_bbans(prior_precision,
beta_binomial_precision,
q_precision, generative_model,
recognition_model, number_buckets,
latent_shape)
#
# Get images to compress
#
images = datasets.get_binarized_MNIST(rng, False)[:image_count]
images = [np.atleast_2d(image.flatten()) for image in images]
original_length = 32 * image_count * 784 #using a float32 per pixel. (not binary)
# generate some random bits for the
random_bits = rng.randint(low=0,
high=((1 << 32) - 1),
size=random_count,
dtype=np.uint32) # some amount of bits
ans.from_array(random_bits)
print("Encoding...")
for i in range(0, len(images)):
image = images[i]
append(ans, image)
compressed_length = 32 * len(ans.to_array())
print("Decoding...")
for i in range(len(images) - 1, -1, -1):
image = pop(ans)
image = image.detach().numpy()
original_image = images[i][0]
assert all(original_image == image)
# this verifies that all bits from the images have been removed and the ans state is restored
assert all(ans.to_array() == random_bits)
print("Decoded all images successfully")
result = Result()
result.exp_name = 'Compression of MNIST dataset using BBANS with Beta Binomial distribution'
result.method_name = 'BBANS using VAE with Beta Binomial latent variables'
result.path_to_model = 'OriginalParameters/torch_vae_beta_binomial_params'
result.image_count = image_count
result.latent_precision = beta_binomial_precision
result.prior_precision = prior_precision
result.posterior_precision = q_precision
result.hidden_size = hidden_dim
result.latent_size = latent_dim
result.image_shape = (28, 28)
result.random_bits_count = random_count * 32
result.seed = seed
result.original_length = original_length
result.compressed_length = compressed_length
result.encode_success = True
result.decode_success = True
result.to_file(result_path)
def original_bernoulli_example(number_images, result_path: str):
"""
Test the binary mnist on VAE BBANS
"""
seed = 0
rng = RandomState(seed)
image_count = number_images
prior_precision = 8
bernoulli_precision = 12
q_precision = 14
random_count = 12
latent_dim = 40
hidden_dim = 100
ans = ANSCoder()
result = Result()
model = bin_vae_original.BinaryVAE(hidden_dim=hidden_dim,
latent_dim=latent_dim)
model.load_state_dict(
torch.load('OriginalParameters/torch_binary_vae_params_new'))
generative_model = torch_to_numpy_function(model.decode)
recognition_model = torch_to_numpy_function(model.encode)
latent_shape = (latent_dim, )
append, pop = build_bernoulli_bbans(prior_precision, bernoulli_precision,
q_precision, generative_model,
recognition_model, latent_shape)
# Get images to compress
images = datasets.get_binarized_MNIST(rng, False)[:image_count]
images = [image.flatten() for image in images]
original_length = 32 * len(images) * len(
images[0]
) #using a float32 per pixel. Could be optimized to 8 bits per pixel.
# generate some random bits
random_bits = rng.randint(low=0,
high=((1 << 32) - 1),
size=random_count,
dtype=np.uint32)
ans.from_array(random_bits)
print("Encoding...")
for i in range(0, len(images)):
image = images[i]
append(ans, image)
compressed_length = 32 * len(ans.to_array())
print("Decoding...")
for i in range(len(images) - 1, -1, -1):
image = pop(ans)
image = image.detach().numpy()
original_image = images[i]
assert all(original_image == image)
# this verifies that all bits from the images have been removed and the ans state is restored
assert all(ans.to_array() == random_bits)
print("All images decoded successfully!")
result = Result()
result.exp_name = 'Compression of binarized MNIST dataset using BBANS with Bernoulli distribution'
result.method_name = 'BBANS using VAE with Bernoulli latent variables'
result.path_to_model = 'OriginalParameters/torch_binary_vae_params_new'
result.image_count = image_count
result.latent_precision = bernoulli_precision
result.prior_precision = prior_precision
result.posterior_precision = q_precision
result.hidden_size = hidden_dim
result.latent_size = latent_dim
result.image_shape = (28, 28)
result.random_bits_count = random_count * 32
result.seed = seed
result.original_length = original_length
result.compressed_length = compressed_length
result.encode_success = True
result.decode_success = True
result.to_file(result_path)
def append(ans: ANSCoder, symbols):
for symbol, cdf in reversed(list(zip(symbols, cdfs))):
start = cdf(symbol)
frequency = cdf(symbol + 1) - start
ans.append(start, frequency, precision)
def append(ans: ANSCoder, symbols):
for symbol in reversed(symbols):
start = symbol
frequency = 1
ans.append(start, frequency, precision)
return