def _test_bits_estimate_matches_entropy(self, prior_fn, base_prior_fn, scale): em = universal.UniversalIndexedEntropyModel( prior_fn=prior_fn, index_ranges=(10, 10), parameter_fns=dict( loc=lambda i: (i[..., 0] - 5.0) / 10.0, scale=lambda i: scale * tf.exp(-i[..., 1]), ), coding_rank=1, compression=True, tail_mass=1e-5, num_noise_levels=15) num_symbols = 1000000 base_df = base_prior_fn(loc=0.1, scale=scale) x = base_df.sample((1, 1, num_symbols), seed=0) # We predict the distribution correctly. indexes = tf.concat((6 * tf.ones((1, 1, num_symbols, 1)), 0 * tf.ones( (1, 1, num_symbols, 1))), axis=-1) bits = em(x, indexes)[1].numpy()[0, 0] bits_per_symbol = bits / num_symbols df = prior_fn(loc=0.0, scale=scale) # Lets estimate entropy via sampling the distribution. samples = df.sample(num_symbols, seed=0) log_probs = df.log_prob(samples) / tf.math.log(2.0) entropy = -tf.reduce_mean(log_probs) # Should be very close. self.assertAllClose(bits_per_symbol, entropy, rtol=0.001, atol=1e-3)
def test_quantization_noise_is_uniform(self, scale): em = universal.UniversalIndexedEntropyModel( prior_fn=uniform_noise.NoisyNormal, index_ranges=(10, 10), parameter_fns=dict( loc=lambda i: (i[..., 0] - 5.0) / 10.0, scale=lambda i: scale * tf.exp(-i[..., 1]), ), coding_rank=1, compression=True) num_symbols = 10000 # Source distribution is gaussian with stddev `scale`. x = tf.random.stateless_normal((1, 1, num_symbols), stddev=scale, seed=(0, 0)) # We predict the distribution correctly. indexes = tf.concat((5 * tf.ones((1, 1, num_symbols, 1)), 0 * tf.ones( (1, 1, num_symbols, 1))), axis=-1) bitstring = em.compress(x, indexes) x_hat = em.decompress(bitstring, indexes) # Quantization noise should be between -.5 and .5 u = x - x_hat self.assertAllLessEqual(tf.abs(u), 0.5) # Check distribution has right statistics. _, p = scipy.stats.kstest(u, "uniform", (-0.5, 1.0)) self.assertGreater(p, 1e-6)
def test_bitstring_length_matches_estimates(self, training): em = universal.UniversalIndexedEntropyModel( prior_fn=uniform_noise.NoisyNormal, index_ranges=(10, 10), parameter_fns=dict( loc=lambda i: (i[..., 0] - 5.0) / 10.0, scale=lambda i: tf.exp(-i[..., 1]), ), coding_rank=1, compression=True) self.assertEqual(em.coding_rank, 1) self.assertEqual(em._laplace_tail_mass, 0.0) self.assertEqual(em.tail_mass, 2**-8) self.assertEqual(em.range_coder_precision, 12) self.assertEqual(em.dtype, tf.float32) num_symbols = 1000 # Source distribution is gaussian with stddev 1. x = tf.random.stateless_normal((1, 1, num_symbols), seed=(0, 0)) # We predict the distribution correctly (mean = 5-5=0, scale = exp(0)=1. indexes = tf.concat((5 * tf.ones((1, 1, num_symbols, 1)), 0 * tf.ones( (1, 1, num_symbols, 1))), axis=-1) x_perturbed, bits_estimate = em(x, indexes, training=training) bitstring = em.compress(x, indexes) x_decoded = em.decompress(bitstring, indexes) bitstring_bits = tf.reshape( [len(b) * 8 for b in bitstring.numpy().flatten()], bitstring.shape) # Max error 1% and 1 byte. self.assertAllClose(bits_estimate, bitstring_bits, atol=8, 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_bitstring_length_matches_entropy(self, prior_fn, base_prior_fn, scale): em = universal.UniversalIndexedEntropyModel( prior_fn=prior_fn, index_ranges=(10, 10), parameter_fns=dict( loc=lambda i: (i[..., 0] - 5.0) / 10.0, scale=lambda i: scale * tf.exp(-i[..., 1]), ), coding_rank=1, compression=True, tail_mass=1e-5, num_noise_levels=15) num_symbols = 10000 base_df = base_prior_fn(loc=0.0, scale=scale) x = base_df.sample((1, 1, num_symbols), seed=0) # We predict the distribution correctly. indexes = tf.concat((5 * tf.ones((1, 1, num_symbols, 1)), 0 * tf.ones( (1, 1, num_symbols, 1))), axis=-1) bitstring = em.compress(x, indexes) bits = len(bitstring.numpy()[0, 0]) * 8 bits_per_symbol = bits / num_symbols df = prior_fn(loc=0.0, scale=scale) # Lets estimate entropy via sampling the distribution. samples = df.sample(num_symbols, seed=0) log_probs = df.log_prob(samples) / tf.math.log(2.0) entropy = -tf.reduce_mean(log_probs) rtol = 0.01 # Maximum relative error 1%. atol = 16 / num_symbols # Maximum 2 bytes absolute error. self.assertLessEqual(bits_per_symbol, entropy * (1 + rtol) + atol)
def test_accurate_predictions_give_small_bitstring_length(self): # If we can perfectly predict locations with a very small scale, the # bitstring_length should be very small. em = universal.UniversalIndexedEntropyModel( prior_fn=uniform_noise.NoisyNormal, index_ranges=(10, 10), parameter_fns=dict( loc=lambda i: i[..., 0] / 3.0, scale=lambda i: tf.exp(-i[..., 1] - 20), # Very small scale. ), coding_rank=1, compression=True) num_symbols = 1000 # Random indices in the valid index ranges. indexes = tf.cast( tf.random.stateless_uniform((3, num_symbols, 16, 2), minval=0.0, maxval=10.0, seed=(0, 0)), tf.int32) prior = em._make_prior(indexes) x = prior.base.sample((), seed=0) y = prior.base.sample((), seed=0) self.assertAllClose(x, y) bitstring = em.compress(x, indexes) bitstring_bits = tf.reshape( [len(b) * 8 for b in bitstring.numpy().flatten()], bitstring.shape) # Maximum 2 bytes. self.assertAllLessEqual(bitstring_bits, 16.0) x_decoded = em.decompress(bitstring, indexes) # Maximum error is .5 self.assertAllLessEqual(tf.abs(x - x_decoded), 0.5)
def test_cannot_instantiate_one_dimensional(self): with self.assertRaises(ValueError): universal.UniversalIndexedEntropyModel( uniform_noise.NoisyNormal, coding_rank=1, index_ranges=64, parameter_fns=dict(loc=lambda _: 0, scale=lambda i: tf.exp(i / 8 - 5)))
def test_expected_grads_or_not_gives_same_bits(self): x = tf.random.stateless_normal((3, 10000, 16), seed=(0, 0)) indexes = tf.cast( 10 * tf.random.stateless_uniform((3, 10000, 16, 3), seed=(0, 0)), tf.int32) em_expected = universal.UniversalIndexedEntropyModel( uniform_noise.NoisyLogisticMixture, index_ranges=(10, 10, 5), parameter_fns=dict( loc=lambda i: i[..., 0:2] - 5, scale=lambda _: 1, weight=lambda i: tf.nn.softmax((i[..., 2:3] - 2) * [-1, 1]), ), coding_rank=2, expected_grads=True) self.assertTrue(em_expected._expected_grads) x_hat, bits_expected = em_expected(x, indexes) # Quantization noise should be between -.5 and .5 u = x - x_hat self.assertAllLessEqual(tf.abs(u), 0.5) em_not_expected = universal.UniversalIndexedEntropyModel( uniform_noise.NoisyLogisticMixture, index_ranges=(10, 10, 5), parameter_fns=dict( loc=lambda i: i[..., 0:2] - 5, scale=lambda _: 1, weight=lambda i: tf.nn.softmax((i[..., 2:3] - 2) * [-1, 1]), ), coding_rank=2, expected_grads=False) self.assertFalse(em_not_expected._expected_grads) x_hat, bits_not_expected = em_not_expected(x, indexes) # Quantization noise should be between -.5 and .5 u = x - x_hat self.assertAllLessEqual(tf.abs(u), 0.5) self.assertAllClose(bits_not_expected, bits_expected, rtol=0.001)
def test_can_instantiate_n_dimensional(self): em = universal.UniversalIndexedEntropyModel( uniform_noise.NoisyLogisticMixture, index_ranges=(10, 10, 5), parameter_fns=dict( loc=lambda i: i[..., 0:2] - 5, scale=lambda _: 1, weight=lambda i: tf.nn.softmax((i[..., 2:3] - 2) * [-1, 1]), ), coding_rank=1, ) self.assertEqual(em.coding_rank, 1) self.assertEqual(em._laplace_tail_mass, 0.0) self.assertEqual(em.tail_mass, 2**-8) self.assertEqual(em.dtype, tf.float32)
def test_can_instantiate_and_compress_n_dimensional(self): em = universal.UniversalIndexedEntropyModel( uniform_noise.NoisyLogisticMixture, index_ranges=(10, 10, 5), parameter_fns=dict( loc=lambda i: i[..., 0:2] - 5, scale=lambda _: 1, weight=lambda i: tf.nn.softmax((i[..., 2:3] - 2) * [-1, 1]), ), coding_rank=1, compression=True) x = tf.random.stateless_normal((3, 8, 16), seed=(0, 0)) indexes = tf.cast( 10 * tf.random.stateless_uniform((3, 8, 16, 3), seed=(0, 0)), tf.int32) em(x, indexes) bitstring = em.compress(x, indexes) x_hat = em.decompress(bitstring, indexes) # Quantization noise should be between -.5 and .5 u = x - x_hat self.assertAllLessEqual(tf.abs(u), 0.5)
def test_expected_grads_gives_gradients(self): x = tf.random.stateless_normal((3, 10000, 16), seed=(0, 0)) indexes = tf.cast(10 * tf.random.stateless_uniform( (3, 10000, 16, 3), seed=(0, 0)), tf.float32) em = universal.UniversalIndexedEntropyModel( uniform_noise.NoisyLogisticMixture, index_ranges=(10, 10, 5), parameter_fns=dict( loc=lambda i: i[..., 0:2] - 5, scale=lambda _: 1, weight=lambda i: tf.nn.softmax((i[..., 2:3] - 2) * [-1, 1]), ), coding_rank=2, expected_grads=True) self.assertTrue(em._expected_grads) with tf.GradientTape(persistent=True) as g: g.watch((x, indexes)) x2, bits = em(x, indexes) self.assertIsInstance(g.gradient(x2, x), tf.Tensor) self.assertIsInstance(g.gradient(bits, x), tf.Tensor) self.assertIsInstance(g.gradient(bits, indexes), tf.Tensor) self.assertIsNone(g.gradient(x2, indexes))