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))
