Add a universally quantized entropy model to tensorflow_compression.

PiperOrigin-RevId: 341791361
Change-Id: Ib80a9ea38f43853d84c76985a8eac0d918d705c7

Author: relational

tensorflow_compression/__init__.py

@@ -29,6 +29,7 @@
 from tensorflow_compression.python.distributions.uniform_noise import *
 from tensorflow_compression.python.entropy_models.continuous_batched import *
 from tensorflow_compression.python.entropy_models.continuous_indexed import *
+from tensorflow_compression.python.entropy_models.universal import *
 from tensorflow_compression.python.layers.entropy_models import *
 from tensorflow_compression.python.layers.gdn import *
 from tensorflow_compression.python.layers.initializers import *

tensorflow_compression/python/distributions/BUILD

@@ -88,6 +88,7 @@ py_test(
     deps = [
         ":deep_factorized",
         ":round_adapters",
+        "//tensorflow_compression/python/ops:soft_round_ops",
     ],
 )
 

tensorflow_compression/python/distributions/helpers.py

@@ -104,6 +104,9 @@ def quantization_offset(distribution):
   these are implemented, it falls back on quantizing to integer values (i.e.,
   an offset of zero).
 
+  Note the offset is always in the range [-.5, .5] as it is assumed to be
+  combined with a round quantizer.
+
   Arguments:
     distribution: A `tfp.distributions.Distribution` object.
 
@@ -125,7 +128,7 @@ def quantization_offset(distribution):
           offset = distribution.mean()
         except NotImplementedError:
           offset = tf.constant(0, dtype=distribution.dtype)
-  return tf.stop_gradient(offset)
+  return tf.stop_gradient(offset - tf.round(offset))
 
 
 def lower_tail(distribution, tail_mass):

tensorflow_compression/python/distributions/helpers_test.py

@@ -23,29 +23,29 @@
 
 class HelpersTest(tf.test.TestCase):
 
-  def test_cauchy_quantizes_to_mode(self):
-    dist = tfp.distributions.Cauchy(loc=1.5, scale=3.)
-    self.assertEqual(helpers.quantization_offset(dist), 1.5)
+  def test_cauchy_quantizes_to_mode_decimal_part(self):
+    dist = tfp.distributions.Cauchy(loc=1.4, scale=3.)
+    self.assertAllClose(helpers.quantization_offset(dist), 0.4)
 
-  def test_gamma_quantizes_to_mode(self):
+  def test_gamma_quantizes_to_mode_decimal_part(self):
     dist = tfp.distributions.Gamma(concentration=5., rate=1.)
-    self.assertEqual(helpers.quantization_offset(dist), 4.)
+    self.assertEqual(helpers.quantization_offset(dist), 0.)
 
-  def test_laplace_quantizes_to_mode(self):
+  def test_laplace_quantizes_to_mode_decimal_part(self):
     dist = tfp.distributions.Laplace(loc=-2., scale=5.)
-    self.assertEqual(helpers.quantization_offset(dist), -2.)
+    self.assertEqual(helpers.quantization_offset(dist), 0.)
 
-  def test_logistic_quantizes_to_mode(self):
+  def test_logistic_quantizes_to_mode_decimal_part(self):
     dist = tfp.distributions.Logistic(loc=-3., scale=1.)
-    self.assertEqual(helpers.quantization_offset(dist), -3.)
+    self.assertEqual(helpers.quantization_offset(dist), 0.)
 
-  def test_lognormal_quantizes_to_mode(self):
+  def test_lognormal_quantizes_to_mode_decimal_part(self):
     dist = tfp.distributions.LogNormal(loc=4., scale=1.)
-    self.assertEqual(helpers.quantization_offset(dist), tf.exp(3.))
+    self.assertAllClose(helpers.quantization_offset(dist), tf.exp(3.)-20.0)
 
-  def test_normal_quantizes_to_mode(self):
+  def test_normal_quantizes_to_mode_decimal_part(self):
     dist = tfp.distributions.Normal(loc=3., scale=5.)
-    self.assertEqual(helpers.quantization_offset(dist), 3.)
+    self.assertEqual(helpers.quantization_offset(dist), 0.)
 
   def test_cauchy_tails_are_in_order(self):
     dist = tfp.distributions.Cauchy(loc=1.5, scale=3.)

tensorflow_compression/python/distributions/round_adapters_test.py

@@ -21,6 +21,7 @@
 
 from tensorflow_compression.python.distributions import deep_factorized
 from tensorflow_compression.python.distributions import round_adapters
+from tensorflow_compression.python.ops import soft_round_ops
 
 
 def _test_log_prob_gradient_is_bounded(self, dist_cls, values, params=()):
@@ -49,19 +50,21 @@ class AdaptersTest(tf.test.TestCase, parameterized.TestCase):
        deep_factorized.DeepFactorized, 0.0),
       ("softround_logistic",
        lambda d: round_adapters.SoftRoundAdapter(d, alpha=5.0),
-       lambda: tfp.distributions.Logistic(loc=10, scale=1.5), 10.0),
+       lambda: tfp.distributions.Logistic(loc=10.3, scale=1.5),
+       lambda: soft_round_ops.soft_round(0.3, alpha=5.0)),
       ("softround_normal",
-       lambda d: round_adapters.SoftRoundAdapter(d, alpha=5.0),
-       lambda: tfp.distributions.Normal(loc=10, scale=1.5), 10.0),
+       lambda d: round_adapters.SoftRoundAdapter(d, alpha=4.0),
+       lambda: tfp.distributions.Normal(loc=10.4, scale=1.5),
+       lambda: soft_round_ops.soft_round(0.4, alpha=4.0)),
       ("noisysoftround_deepfactorized",
        lambda d: round_adapters.NoisySoftRoundAdapter(d, alpha=5.0),
        deep_factorized.DeepFactorized, 0.0),
       ("noisysoftround_logistic",
        lambda d: round_adapters.NoisySoftRoundAdapter(d, alpha=5.0),
-       lambda: tfp.distributions.Logistic(loc=10, scale=1.5), 10.0),
+       lambda: tfp.distributions.Logistic(loc=10, scale=1.5), 0.0),
       ("noisysoftround_normal",
        lambda d: round_adapters.NoisySoftRoundAdapter(d, alpha=5.0),
-       lambda: tfp.distributions.Normal(loc=10, scale=1.5), 10.0),
+       lambda: tfp.distributions.Normal(loc=10, scale=1.5), 0.0),
       ("round_deepfactorized",
        round_adapters.RoundAdapter,
        lambda: deep_factorized.DeepFactorized(init_scale=1.0), 0.0),
@@ -101,6 +104,10 @@ def test_tails_and_offset(self, adapter, distribution, expected_offset):
 
     self.assertGreater(upper_tail, lower_tail)
     offset = dist._quantization_offset()
+    if not isinstance(expected_offset, float):
+      # We cannot run tf inside the parameterized test declaration, hence
+      # non-float values are wrapped in a lambda.
+      expected_offset = expected_offset()
     self.assertAllClose(offset, expected_offset)
 
   @parameterized.named_parameters(

tensorflow_compression/python/distributions/uniform_noise_test.py

@@ -60,12 +60,12 @@ def test_sampling_works(self):
     self.assertEqual(sample.shape, (5, 4, 2))
 
   def test_tails_and_offset_are_in_order(self):
-    dist = self.dist_cls(loc=10, scale=1.5)
+    dist = self.dist_cls(loc=10.3, scale=1.5)
     offset = helpers.quantization_offset(dist)
     lower_tail = helpers.lower_tail(dist, 2**-8)
     upper_tail = helpers.upper_tail(dist, 2**-8)
-    self.assertGreater(upper_tail, offset)
-    self.assertGreater(offset, lower_tail)
+    self.assertGreater(upper_tail, lower_tail)
+    self.assertAllClose(offset, 0.3)
 
   def test_stats_throw_error(self):
     dist = self.dist_cls(loc=1, scale=2)
@@ -130,12 +130,12 @@ def test_sampling_works(self):
     self.assertEqual(sample.shape, (5, 4, 1))
 
   def test_tails_and_offset_are_in_order(self):
-    dist = self.dist_cls(loc=10, scale=[1.5, 2], weight=[.5, .5])
+    dist = self.dist_cls(loc=[5.4, 8.6], scale=[1.4, 2], weight=[.6, .4])
     offset = helpers.quantization_offset(dist)
     lower_tail = helpers.lower_tail(dist, 2**-8)
     upper_tail = helpers.upper_tail(dist, 2**-8)
-    self.assertGreater(upper_tail, offset)
-    self.assertGreater(offset, lower_tail)
+    self.assertGreater(upper_tail, lower_tail)
+    self.assertAllClose(offset, 0.4)  # Decimal part of the peakiest mode (5.4).
 
   def test_stats_throw_error(self):
     dist = self.dist_cls(loc=[1, 0], scale=2, weight=[.1, .9])

tensorflow_compression/python/entropy_models/BUILD

@@ -12,6 +12,7 @@ py_library(
         ":continuous_base",
         ":continuous_batched",
         ":continuous_indexed",
+        ":universal",
     ],
 )
 
@@ -32,6 +33,7 @@ py_library(
     deps = [
         ":continuous_base",
         "//tensorflow_compression/python/distributions:helpers",
+        "//tensorflow_compression/python/ops:math_ops",
         "//tensorflow_compression/python/ops:range_coding_ops",
     ],
 )
@@ -68,6 +70,30 @@ py_test(
     ],
 )
 
+py_library(
+    name = "universal",
+    srcs = ["universal.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":continuous_batched",
+        ":continuous_indexed",
+        "//tensorflow_compression/python/ops:math_ops",
+    ],
+)
+
+py_test(
+    name = "universal_test",
+    timeout = "long",
+    srcs = ["universal_test.py"],
+    python_version = "PY3",
+    shard_count = 3,
+    deps = [
+        ":universal",
+        "//tensorflow_compression/python/distributions:deep_factorized",
+        "//tensorflow_compression/python/distributions:uniform_noise",
+    ],
+)
+
 filegroup(
     name = "py_src",
     srcs = glob(["*.py"]),