def test_laplace_tail_mass_for_small_inputs(self): prior = deep_factorized.NoisyDeepFactorized(batch_shape=(1, )) em1 = universal.UniversalBatchedEntropyModel(prior, coding_rank=1, compression=True, laplace_tail_mass=1e-3) em2 = universal.UniversalBatchedEntropyModel(prior, coding_rank=1, compression=True) x = tf.linspace(-10.0, 10.0, 50) _, bits1 = em1(x[..., None]) _, bits2 = em2(x[..., None]) self.assertAllClose(bits1, bits2, rtol=0.01, atol=0.05)
def test_bitstring_length_matches_estimates(self, training, prior): priors = { "deep_factorized": deep_factorized.NoisyDeepFactorized(batch_shape=(16, )), "normal": uniform_noise.NoisyNormal(loc=tf.range(16.0), scale=1.0) } prior = priors[prior] em = universal.UniversalBatchedEntropyModel(prior, coding_rank=2, compression=True) num_symbols = 1000 # Source distribution is fixed as gaussian. source = priors["normal"].base x = source.sample((3, num_symbols), seed=0) x_perturbed, bits_estimate = em(x, training=training) bitstring = em.compress(x) x_decoded = em.decompress(bitstring, (num_symbols, )) bitstring_bits = tf.reshape( [len(b) * 8 for b in bitstring.numpy().flatten()], bitstring.shape) # Max error 1% and 2 bytes. self.assertAllClose(bits_estimate, bitstring_bits, atol=16, rtol=0.01) # Quantization noise should be between -.5 and .5 self.assertAllLessEqual(tf.abs(x - x_decoded), 0.5) self.assertAllLessEqual(tf.abs(x - x_perturbed), 0.5)
def test_laplace_tail_mass_for_large_inputs(self): prior = deep_factorized.NoisyDeepFactorized(batch_shape=(1, )) em = universal.UniversalBatchedEntropyModel(prior, coding_rank=1, compression=True, laplace_tail_mass=1e-3) x = tf.convert_to_tensor([1e3, 1e4, 1e5, 1e6, 1e7, 1e8], tf.float32) _, bits = em(x[..., None]) self.assertAllClose(bits, tf.abs(x) / tf.math.log(2.0), rtol=0.01)
def test_can_instantiate_and_compress(self): prior = deep_factorized.NoisyDeepFactorized(batch_shape=(4, 4)) em = universal.UniversalBatchedEntropyModel( prior, coding_rank=3, compression=True) x = tf.random.stateless_normal((3, 8, 4, 4), seed=(0, 0)) bitstring = em.compress(x) em(x) x_hat = em.decompress(bitstring, (8,)) # Quantization noise should be between -.5 and .5 u = x - x_hat self.assertAllLessEqual(tf.abs(u), 0.5)
def test_bitstring_length_matches_entropy_normal(self, scale=1e-8): prior = uniform_noise.NoisyNormal(loc=100 * tf.range(15.0), scale=scale) base_df = prior.base em = universal.UniversalBatchedEntropyModel( prior, coding_rank=2, compression=True) num_samples = 100000 x = base_df.sample(num_samples, seed=0) bitstring = em.compress(x) x_decoded = em.decompress(bitstring, (num_samples,)) bits = len(bitstring.numpy()) * 8 bits_per_sample = bits / num_samples # Quantization noise should be between -.5 and .5 self.assertAllLessEqual(tf.abs(x - x_decoded), 0.5) # Lets estimate entropy via sampling the distribution. samples = prior.sample(num_samples, seed=0) log_probs = prior.log_prob(samples) / tf.math.log(2.0) entropy_bits = -tf.reduce_sum(log_probs) rtol = 0.01 # Maximum relative error 1%. atol = 16 # Maximum 2 bytes absolute error. self.assertLessEqual(bits_per_sample, entropy_bits * (1 + rtol) + atol)
def test_expected_grads_gives_gradients(self): priors = { "deep_factorized": deep_factorized.NoisyDeepFactorized(batch_shape=(16,)), "normal": uniform_noise.NoisyNormal(loc=tf.range(16.0), scale=1.0) } prior = priors["deep_factorized"] em = universal.UniversalBatchedEntropyModel( prior, coding_rank=2, compression=True, expected_grads=True) self.assertTrue(em._expected_grads) num_symbols = 1000 # Source distribution is fixed as gaussian. source = priors["normal"].base x = source.sample((3, num_symbols), seed=0) with tf.GradientTape(persistent=True) as g: g.watch(x) x2, bits = em(x, training=True) self.assertIsInstance(g.gradient(x2, x), tf.Tensor) self.assertIsInstance(g.gradient(bits, x), tf.Tensor) for variable in em.trainable_variables: self.assertIsInstance(g.gradient(bits, variable), tf.Tensor)