def test_tails_and_offset_are_in_order(self): df = deep_factorized.NoisyDeepFactorized() offset = helpers.quantization_offset(df) lower_tail = helpers.lower_tail(df, 2**-8) upper_tail = helpers.upper_tail(df, 2**-8) self.assertGreater(upper_tail, offset) self.assertGreater(offset, lower_tail)
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_can_instantiate_and_run_scalar(self): df = deep_factorized.NoisyDeepFactorized(num_filters=(2, 3, 4)) self.assertEqual(df.batch_shape, ()) self.assertEqual(df.event_shape, ()) self.assertEqual(df.base.num_filters, (2, 3, 4)) self.assertEqual(df.base.init_scale, 10) x = tf.random.normal((10, )) df.prob(x)
def test_variables_receive_gradients(self): df = deep_factorized.NoisyDeepFactorized() with tf.GradientTape() as tape: x = tf.random.normal([20]) loss = -tf.reduce_mean(df.log_prob(x)) grads = tape.gradient(loss, df.trainable_variables) self.assertLen(grads, 8) self.assertNotIn(None, grads)
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_logistic_is_special_case(self): # With no hidden units, the density should collapse to a logistic # distribution convolved with a standard uniform distribution. df = deep_factorized.NoisyDeepFactorized(num_filters=(), init_scale=1) logistic = tfp.distributions.Logistic(loc=-df.base._biases[0][0, 0], scale=1.) x = tf.linspace(-5., 5., 20) prob_df = df.prob(x) prob_log = logistic.cdf(x + .5) - logistic.cdf(x - .5) self.assertAllClose(prob_df, prob_log)
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_stats_throw_error(self): df = deep_factorized.NoisyDeepFactorized() with self.assertRaises(NotImplementedError): df.mode() with self.assertRaises(NotImplementedError): df.mean() with self.assertRaises(NotImplementedError): df.quantile(.5) with self.assertRaises(NotImplementedError): df.survival_function(.5) with self.assertRaises(NotImplementedError): df.sample()
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_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)
def test_quantization_offset_is_zero(self): df = deep_factorized.NoisyDeepFactorized() self.assertEqual(helpers.quantization_offset(df), 0)
def test_uniform_is_special_case(self): # With the scale parameter going to zero, the density should approach a # unit-width uniform distribution. df = deep_factorized.NoisyDeepFactorized(init_scale=1e-3) x = tf.linspace(-1., 1., 10) self.assertAllClose(df.prob(x), [0, 0, 0, 1, 1, 1, 1, 0, 0, 0])
def test_tails_are_in_order(self): df = deep_factorized.NoisyDeepFactorized() lower_tail = helpers.lower_tail(df, 2**-8) upper_tail = helpers.upper_tail(df, 2**-8) self.assertGreater(upper_tail, lower_tail)
def test_noisy_deep_factorized_tails_are_in_order(self): dist = deep_factorized.NoisyDeepFactorized(batch_shape=[10]) self.assertAllGreater( helpers.upper_tail(dist, 2**-8) - helpers.lower_tail(dist, 2**-8), 0)