Adds no_variables option to entropy models.

PiperOrigin-RevId: 332352604
Change-Id: I1fd01aeeda87e6b27d22431e0a7c9467eb2e5fb1

Author: Johannes Ballé

tensorflow_compression/python/entropy_models/continuous_base.py

@@ -39,7 +39,7 @@ class ContinuousEntropyModelBase(tf.Module, metaclass=abc.ABCMeta):
   @abc.abstractmethod
   def __init__(self, prior, coding_rank, compression=False,
                likelihood_bound=1e-9, tail_mass=2**-8,
-               range_coder_precision=12):
+               range_coder_precision=12, no_variables=False):
     """Initializer.
 
     Arguments:
@@ -60,6 +60,8 @@ def __init__(self, prior, coding_rank, compression=False,
       tail_mass: Float. Approximate probability mass which is range encoded with
         less precision, by using a Golomb-like code.
       range_coder_precision: Integer. Precision passed to the range coding op.
+      no_variables: Boolean. If True, creates range coding tables as `Tensor`s
+        rather than `Variable`s.
 
     Raises:
       RuntimeError: when attempting to instantiate an entropy model with
@@ -77,6 +79,7 @@ def __init__(self, prior, coding_rank, compression=False,
       self._likelihood_bound = float(likelihood_bound)
       self._tail_mass = float(tail_mass)
       self._range_coder_precision = int(range_coder_precision)
+      self._no_variables = bool(no_variables)
       if self.compression:
         self._build_tables(prior)
 
@@ -103,17 +106,17 @@ def _check_compression(self):
   @property
   def cdf(self):
     self._check_compression()
-    return tf.identity(self._cdf)
+    return tf.convert_to_tensor(self._cdf)
 
   @property
   def cdf_offset(self):
     self._check_compression()
-    return tf.identity(self._cdf_offset)
+    return tf.convert_to_tensor(self._cdf_offset)
 
   @property
   def cdf_length(self):
     self._check_compression()
-    return tf.identity(self._cdf_length)
+    return tf.convert_to_tensor(self._cdf_length)
 
   @property
   def dtype(self):
@@ -155,6 +158,11 @@ def range_coder_precision(self):
     """Precision passed to range coding op."""
     return self._range_coder_precision
 
+  @property
+  def no_variables(self):
+    """Whether range coding tables are created as `Tensor`s or `Variable`s."""
+    return self._no_variables
+
   @tf.custom_gradient
   def _quantize_no_offset(self, inputs):
     return tf.round(inputs), lambda x: x
@@ -247,11 +255,16 @@ def loop_body(args):
       cdf = tf.map_fn(
           loop_body, (pmf, pmf_length), dtype=tf.int32, name="pmf_to_cdf")
 
-    self._cdf = tf.Variable(cdf, trainable=False, name="cdf")
-    self._cdf_offset = tf.Variable(
-        cdf_offset, trainable=False, name="cdf_offset")
-    self._cdf_length = tf.Variable(
-        cdf_length, trainable=False, name="cdf_length")
+    if self.no_variables:
+      self._cdf = cdf
+      self._cdf_offset = cdf_offset
+      self._cdf_length = cdf_length
+    else:
+      self._cdf = tf.Variable(cdf, trainable=False, name="cdf")
+      self._cdf_offset = tf.Variable(
+          cdf_offset, trainable=False, name="cdf_offset")
+      self._cdf_length = tf.Variable(
+          cdf_length, trainable=False, name="cdf_length")
 
   @abc.abstractmethod
   def get_config(self):
@@ -264,10 +277,10 @@ def get_config(self):
       NotImplementedError: on attempting to call this method on an entropy model
         with `compression=False`.
     """
-    if not self.compression:
+    if self.no_variables or not self.compression:
       raise NotImplementedError(
-          "Serializing entropy models with `compression=False` is currently "
-          "not supported.")
+          "Serializing entropy models with `compression=False` or "
+          "`no_variables=True` is currently not supported.")
     return dict(
         dtype=self._dtype.name,
         prior_shape=self._prior_shape,
@@ -304,6 +317,7 @@ def from_config(cls, config):
       self._likelihood_bound = float(config["likelihood_bound"])
       self._tail_mass = float(config["tail_mass"])
       self._range_coder_precision = int(config["range_coder_precision"])
+      self._no_variables = False
 
       prior_size = functools.reduce(lambda x, y: x * y, self.prior_shape, 1)
       cdf_width = int(config["cdf_width"])

tensorflow_compression/python/entropy_models/continuous_batched.py

@@ -75,7 +75,7 @@ class ContinuousBatchedEntropyModel(continuous_base.ContinuousEntropyModelBase):
 
   def __init__(self, prior, coding_rank, compression=False,
                likelihood_bound=1e-9, tail_mass=2**-8,
-               range_coder_precision=12):
+               range_coder_precision=12, no_variables=False):
     """Initializer.
 
     Arguments:
@@ -98,6 +98,8 @@ def __init__(self, prior, coding_rank, compression=False,
       tail_mass: Float. Approximate probability mass which is range encoded with
         less precision, by using a Golomb-like code.
       range_coder_precision: Integer. Precision passed to the range coding op.
+      no_variables: Boolean. If True, creates range coding tables as `Tensor`s
+        rather than `Variable`s.
 
     Raises:
       RuntimeError: when attempting to instantiate an entropy model with
@@ -107,27 +109,38 @@ def __init__(self, prior, coding_rank, compression=False,
       raise ValueError(
           "`coding_rank` can't be smaller than batch rank of prior.")
     super().__init__(
-        prior, coding_rank, compression=compression,
-        likelihood_bound=likelihood_bound, tail_mass=tail_mass,
-        range_coder_precision=range_coder_precision)
+        prior=prior,
+        coding_rank=coding_rank,
+        compression=compression,
+        likelihood_bound=likelihood_bound,
+        tail_mass=tail_mass,
+        range_coder_precision=range_coder_precision,
+        no_variables=no_variables,
+    )
 
     quantization_offset = helpers.quantization_offset(prior)
-    if self.compression:
-      # Optimization: if the quantization offset is zero, we don't need to
-      # subtract/add it when quantizing, and we don't need to serialize its
-      # value. Note that this code will only work in eager mode.
-      # TODO(jonycgn): Reconsider if this optimization is worth keeping once
-      # the implementation is stable.
-      if tf.executing_eagerly() and tf.reduce_all(
-          tf.equal(quantization_offset, 0.)):
-        quantization_offset = None
-      else:
-        quantization_offset = tf.broadcast_to(
-            quantization_offset, self.prior_shape_tensor)
+    # Optimization: if the quantization offset is zero, we don't need to
+    # subtract/add it when quantizing, and we don't need to serialize its value.
+    # Note that this code will only work in eager mode.
+    # TODO(jonycgn): Reconsider if this optimization is worth keeping once the
+    # implementation is stable.
+    if tf.executing_eagerly() and tf.reduce_all(
+        tf.equal(quantization_offset, 0.)):
+      quantization_offset = None
+    else:
+      quantization_offset = tf.broadcast_to(
+          quantization_offset, self.prior_shape_tensor)
+      if self.compression and not self.no_variables:
         quantization_offset = tf.Variable(
             quantization_offset, trainable=False, name="quantization_offset")
     self._quantization_offset = quantization_offset
 
+  @property
+  def quantization_offset(self):
+    if self._quantization_offset is None:
+      return None
+    return tf.convert_to_tensor(self._quantization_offset)
+
   def _compute_indexes(self, broadcast_shape):
     # TODO(jonycgn, ssjhv): Investigate broadcasting in range coding op.
     prior_size = functools.reduce(lambda x, y: x * y, self.prior_shape, 1)
@@ -187,7 +200,7 @@ def quantize(self, bottleneck):
     Returns:
       A `tf.Tensor` containing the quantized values.
     """
-    return self._quantize(bottleneck, self._quantization_offset)
+    return self._quantize(bottleneck, self.quantization_offset)
 
   @tf.Module.with_name_scope
   def compress(self, bottleneck):
@@ -220,8 +233,8 @@ def compress(self, bottleneck):
         :self.coding_rank - len(self.prior_shape)]
 
     indexes = self._compute_indexes(broadcast_shape)
-    if self._quantization_offset is not None:
-      bottleneck -= self._quantization_offset
+    if self.quantization_offset is not None:
+      bottleneck -= self.quantization_offset
     symbols = tf.cast(tf.round(bottleneck), tf.int32)
     symbols = tf.reshape(symbols, tf.concat([[-1], coding_shape], 0))
 
@@ -287,8 +300,8 @@ def loop_body(string):
 
     symbols = tf.reshape(symbols, symbols_shape)
     outputs = tf.cast(symbols, self.dtype)
-    if self._quantization_offset is not None:
-      outputs += self._quantization_offset
+    if self.quantization_offset is not None:
+      outputs += self.quantization_offset
     return outputs
 
   def get_config(self):

tensorflow_compression/python/entropy_models/continuous_indexed.py

@@ -129,7 +129,7 @@ class ContinuousIndexedEntropyModel(continuous_base.ContinuousEntropyModelBase):
   def __init__(self, prior_fn, index_ranges, parameter_fns, coding_rank,
                compression=False, channel_axis=-1, dtype=tf.float32,
                likelihood_bound=1e-9, tail_mass=2**-8,
-               range_coder_precision=12):
+               range_coder_precision=12, no_variables=False):
     """Initializer.
 
     Arguments:
@@ -170,6 +170,8 @@ def __init__(self, prior_fn, index_ranges, parameter_fns, coding_rank,
       tail_mass: Float. Approximate probability mass which is range encoded with
         less precision, by using a Golomb-like code.
       range_coder_precision: Integer. Precision passed to the range coding op.
+      no_variables: Boolean. If True, creates range coding tables as `Tensor`s
+        rather than `Variable`s.
 
     Raises:
       RuntimeError: when attempting to instantiate an entropy model with
@@ -204,9 +206,14 @@ def __init__(self, prior_fn, index_ranges, parameter_fns, coding_rank,
     prior = self.prior_fn(**parameters)  # pylint:disable=not-callable
 
     super().__init__(
-        prior, coding_rank, compression=compression,
-        likelihood_bound=likelihood_bound, tail_mass=tail_mass,
-        range_coder_precision=range_coder_precision)
+        prior=prior,
+        coding_rank=coding_rank,
+        compression=compression,
+        likelihood_bound=likelihood_bound,
+        tail_mass=tail_mass,
+        range_coder_precision=range_coder_precision,
+        no_variables=no_variables,
+    )
 
   @property
   def index_ranges(self):
@@ -433,7 +440,7 @@ class LocationScaleIndexedEntropyModel(ContinuousIndexedEntropyModel):
 
   def __init__(self, prior_fn, num_scales, scale_fn, coding_rank,
                compression=False, dtype=tf.float32, likelihood_bound=1e-9,
-               tail_mass=2**-8, range_coder_precision=12):
+               tail_mass=2**-8, range_coder_precision=12, no_variables=False):
     """Initializer.
 
     Arguments:
@@ -464,6 +471,8 @@ def __init__(self, prior_fn, num_scales, scale_fn, coding_rank,
       tail_mass: Float. Approximate probability mass which is range encoded with
         less precision, by using a Golomb-like code.
       range_coder_precision: Integer. Precision passed to the range coding op.
+      no_variables: Boolean. If True, creates range coding tables as `Tensor`s
+        rather than `Variable`s.
     """
     num_scales = int(num_scales)
     super().__init__(
@@ -479,6 +488,7 @@ def __init__(self, prior_fn, num_scales, scale_fn, coding_rank,
         likelihood_bound=likelihood_bound,
         tail_mass=tail_mass,
         range_coder_precision=range_coder_precision,
+        no_variables=no_variables,
     )
 
   @tf.Module.with_name_scope