Refactors and updates models to TF2 API/Keras.

Adds support for training/validating on CLIC dataset.

Simplifies PackedTensors, and fixes two bugs with serialization of activation
functions and layers that have not been built.

PiperOrigin-RevId: 361232886
Change-Id: I1ebffe4e89fcfeb899994f0bd9931fb3ac8641b4

Author: Johannes Ballé

models/bls2017.py

@@ -97,7 +97,7 @@ def __init__(self,
         `compress()` and `decompress()` will be built on instantiation. If set
         to `False`, these two methods will not be accessible.
       laplace_tail_mass: Float. If positive, will augment the prior with a
-        laplace mixture for training stability. (experimental)
+        Laplace mixture for training stability. (experimental)
       expected_grads: If True, will use analytical expected gradients during
         backpropagation w.r.t. additive uniform noise.
       tail_mass: Float. Approximate probability mass which is range encoded with
@@ -164,7 +164,6 @@ def _compute_indexes_and_offset(self, broadcast_shape):
   def __call__(self, bottleneck, training=True):
     """Perturbs a tensor with (quantization) noise and estimates bitcost.
 
-
     Args:
       bottleneck: `tf.Tensor` containing the data to be compressed. Must have at
         least `self.coding_rank` dimensions, and the innermost dimensions must
@@ -280,9 +279,8 @@ def decompress(self, strings, broadcast_shape):
     Args:
       strings: `tf.Tensor` containing the compressed bit strings.
       broadcast_shape: Iterable of ints. The part of the output tensor shape
-        between the shape of `strings` on the left and
-        `self.prior_shape` on the right. This must match the shape
-        of the input to `compress()`.
+        between the shape of `strings` on the left and `self.prior_shape` on the
+        right. This must match the shape of the input to `compress()`.
 
     Returns:
       A `tf.Tensor` of shape `strings.shape + broadcast_shape +

models/bmshj2018.py

@@ -423,14 +423,16 @@ def get_config(self) -> Dict[str, Any]:
 
     # Since alpha and epsilon are scalar, allow fixed values to be serialized.
     def try_serialize(parameter, name):
+      if parameter is None:
+        return None
       try:
         return tf.keras.utils.serialize_keras_object(parameter)
       except (ValueError, TypeError):  # Should throw TypeError, but doesn't...
         try:
           return float(parameter)
         except TypeError:
           raise TypeError(
-              f"Can't serialize {name} of type '{type(parameter)}'.")
+              f"Can't serialize {name} of type {type(parameter)}.")
 
     alpha_parameter = try_serialize(
         self.alpha_parameter, "alpha_parameter")

models/ms2020.py

@@ -146,7 +146,8 @@ def test_variables_receive_gradients(self):
     weight_shapes = [tuple(w.shape) for w in layer.trainable_weights]
     self.assertSameElements(grad_shapes, weight_shapes)
 
-  def test_can_be_saved_within_functional_model(self):
+  @parameterized.parameters(False, True)
+  def test_can_be_saved_within_functional_model(self, build):
     inputs = tf.keras.Input(shape=(5,))
     outputs = gdn.GDN()(inputs)
     model = tf.keras.Model(inputs=inputs, outputs=outputs)
@@ -161,12 +162,13 @@ def test_can_be_saved_within_functional_model(self):
       self.assertIsInstance(layer.epsilon_parameter, tf.Tensor)
       self.assertEmpty(layer.epsilon_parameter.shape)
 
-    x = tf.random.uniform((5, 5), dtype=tf.float32)
-    y = model(x)
-    weight_names = [w.name for w in model.weights]
+    if build:
+      x = tf.random.uniform((5, 5), dtype=tf.float32)
+      y = model(x)
+      weight_names = [w.name for w in model.weights]
 
     tempdir = self.create_tempdir()
-    model_path = os.path.join(tempdir.full_path, "model")
+    model_path = os.path.join(tempdir, "model")
     # This should force the model to be reconstructed via configs.
     model.save(model_path, save_traces=False)
 
@@ -182,11 +184,12 @@ def test_can_be_saved_within_functional_model(self):
       self.assertIsInstance(layer.epsilon_parameter, tf.Tensor)
       self.assertEmpty(layer.epsilon_parameter.shape)
 
-    with self.subTest(name="model_outputs_identical"):
-      self.assertAllEqual(model(x), y)
+    if build:
+      with self.subTest(name="model_outputs_identical"):
+        self.assertAllEqual(model(x), y)
 
-    with self.subTest(name="model_weights_identical"):
-      self.assertSameElements(weight_names, [w.name for w in model.weights])
+      with self.subTest(name="model_weights_identical"):
+        self.assertSameElements(weight_names, [w.name for w in model.weights])
 
 
 if __name__ == "__main__":

tensorflow_compression/python/entropy_models/continuous_batched.py

@@ -439,7 +439,7 @@ def activation(self) -> Optional[Callable[[Any], tf.Tensor]]:
   @activation.setter
   def activation(self, value):
     self._check_not_built()
-    self._activation = value
+    self._activation = tf.keras.activations.get(value)
 
   @property
   def use_bias(self) -> bool:
@@ -938,7 +938,7 @@ def call(self, inputs) -> tf.Tensor:
 
     # Finally, pass through activation function if requested.
     if self.activation is not None:
-      outputs = self.activation(outputs)  # pylint:disable=not-callable
+      outputs = self.activation(outputs)
 
     return outputs
 
@@ -979,13 +979,15 @@ def get_config(self) -> Dict[str, Any]:
     # Special-case variables, which can't be serialized but are handled by
     # get_weights()/set_weights().
     def try_serialize(parameter, name):
+      if isinstance(parameter, str):
+        return parameter
       try:
         return tf.keras.utils.serialize_keras_object(parameter)
       except (ValueError, TypeError):  # Should throw TypeError, but doesn't...
         if isinstance(parameter, tf.Variable):
           return "variable"
         raise TypeError(
-            f"Can't serialize {name} of type '{type(parameter)}'.")
+            f"Can't serialize {name} of type {type(parameter)}.")
 
     kernel_parameter = try_serialize(self.kernel_parameter, "kernel")
     bias_parameter = try_serialize(self.bias_parameter, "bias")
@@ -1000,7 +1002,7 @@ def try_serialize(parameter, name):
         extra_pad_end=self.extra_pad_end,
         channel_separable=self.channel_separable,
         data_format=self.data_format,
-        activation=self.activation,
+        activation=tf.keras.activations.serialize(self.activation),
         use_bias=self.use_bias,
         use_explicit=self.use_explicit,
         kernel_parameter=kernel_parameter,

tensorflow_compression/python/layers/gdn.py

@@ -15,6 +15,7 @@
 """Tests of signal processing convolution layers."""
 
 import os
+from absl.testing import parameterized
 import numpy as np
 import scipy.signal
 import tensorflow as tf
@@ -23,7 +24,7 @@
 from tensorflow_compression.python.layers import signal_conv
 
 
-class SignalConvTest(tf.test.TestCase):
+class SignalConvTest(tf.test.TestCase, parameterized.TestCase):
 
   def test_invalid_data_format_raises_error(self):
     with self.assertRaises(ValueError):
@@ -112,9 +113,11 @@ def test_variables_receive_gradients(self):
     weight_shapes = [tuple(w.shape) for w in layer.trainable_weights]
     self.assertSameElements(grad_shapes, weight_shapes)
 
-  def test_can_be_saved_within_functional_model(self):
+  @parameterized.parameters(False, True)
+  def test_can_be_saved_within_functional_model(self, build):
     inputs = tf.keras.Input(shape=(None, 2))
-    outputs = signal_conv.SignalConv1D(1, 3, use_bias=True)(inputs)
+    outputs = signal_conv.SignalConv1D(
+        1, 3, use_bias=True, activation=tf.nn.relu)(inputs)
     model = tf.keras.Model(inputs=inputs, outputs=outputs)
     layer = model.get_layer("signal_conv1d")
 
@@ -123,12 +126,13 @@ def test_can_be_saved_within_functional_model(self):
       self.assertIsInstance(layer.kernel_parameter, parameters.RDFTParameter)
       self.assertIsInstance(layer.bias_parameter, tf.Variable)
 
-    x = tf.random.uniform((1, 5, 2), dtype=tf.float32)
-    y = model(x)
-    weight_names = [w.name for w in model.weights]
+    if build:
+      x = tf.random.uniform((1, 5, 2), dtype=tf.float32)
+      y = model(x)
+      weight_names = [w.name for w in model.weights]
 
     tempdir = self.create_tempdir()
-    model_path = os.path.join(tempdir.full_path, "model")
+    model_path = os.path.join(tempdir, "model")
     # This should force the model to be reconstructed via configs.
     model.save(model_path, save_traces=False)
 
@@ -140,11 +144,12 @@ def test_can_be_saved_within_functional_model(self):
       self.assertIsInstance(layer.kernel_parameter, parameters.RDFTParameter)
       self.assertIsInstance(layer.bias_parameter, tf.Variable)
 
-    with self.subTest(name="model_outputs_identical"):
-      self.assertAllEqual(model(x), y)
+    if build:
+      with self.subTest(name="model_outputs_identical"):
+        self.assertAllEqual(model(x), y)
 
-    with self.subTest(name="model_weights_identical"):
-      self.assertSameElements(weight_names, [w.name for w in model.weights])
+      with self.subTest(name="model_weights_identical"):
+        self.assertSameElements(weight_names, [w.name for w in model.weights])
 
 
 class ConvolutionsTest(tf.test.TestCase):
@@ -353,7 +358,7 @@ def run_or_fail(self, method,
       except:
         msg = []
         for k in sorted(args):
-          msg.append("{}={}".format(k, args[k]))
+          msg.append(f"{k}={args[k]}")
         print("Failed when it shouldn't have: " + ", ".join(msg))
         raise
     else:
@@ -363,7 +368,7 @@ def run_or_fail(self, method,
       except:
         msg = []
         for k in sorted(args):
-          msg.append("{}={}".format(k, args[k]))
+          msg.append(f"{k}={args[k]}")
         print("Did not fail when it should have: " + ", ".join(msg))
         raise
 

tensorflow_compression/python/layers/gdn_test.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 """Packed tensors in bit sequences."""
 
-import numpy as np
 import tensorflow as tf
 
 
@@ -62,50 +61,36 @@ def string(self):
   def string(self, value):
     self._example.ParseFromString(value)
 
-  def pack(self, tensors, arrays):
+  def pack(self, tensors):
     """Packs `Tensor` values into this object."""
-    if len(tensors) != len(arrays):
-      raise ValueError("`tensors` and `arrays` must have same length.")
     i = 1
-    for tensor, array in zip(tensors, arrays):
+    for tensor in tensors:
       feature = self._example.features.feature[chr(i)]
       feature.Clear()
-      if array.ndim != 1:
-        raise RuntimeError("Unexpected tensor rank: {}.".format(array.ndim))
+      if tensor.shape.rank != 1:
+        raise RuntimeError(f"Unexpected tensor rank: {tensor.shape.rank}.")
       if tensor.dtype.is_integer:
-        feature.int64_list.value[:] = array
+        feature.int64_list.value[:] = tensor.numpy()
       elif tensor.dtype == tf.string:
-        feature.bytes_list.value[:] = array
+        feature.bytes_list.value[:] = tensor.numpy()
       else:
-        raise RuntimeError(
-            "Unexpected tensor dtype: '{}'.".format(tensor.dtype))
+        raise RuntimeError(f"Unexpected tensor dtype: '{tensor.dtype}'.")
       i += 1
     # Delete any remaining, previously set arrays.
     while chr(i) in self._example.features.feature:
       del self._example.features.feature[chr(i)]
       i += 1
 
-  # TODO(jonycgn): Remove this function once all models are converted.
-  def unpack(self, tensors):
-    """Unpacks `Tensor` values from this object."""
-    # Check tensor dtype first for a more informative error message.
-    for x in tensors:
-      if not x.dtype.is_integer and x.dtype != tf.string:
-        raise RuntimeError("Unexpected tensor dtype: '{}'.".format(x.dtype))
-
-    # Extact numpy dtypes and call type-based API.
-    np_dtypes = [x.dtype.as_numpy_dtype for x in tensors]
-    return self.unpack_from_np_dtypes(np_dtypes)
-
-  def unpack_from_np_dtypes(self, np_dtypes):
-    """Unpacks values from this object based on numpy dtypes."""
-    arrays = []
-    for i, np_dtype in enumerate(np_dtypes):
+  def unpack(self, dtypes):
+    """Unpacks values from this object based on dtypes."""
+    tensors = []
+    for i, dtype in enumerate(dtypes):
+      dtype = tf.as_dtype(dtype)
       feature = self._example.features.feature[chr(i + 1)]
-      if np.issubdtype(np_dtype, np.integer):
-        arrays.append(np.array(feature.int64_list.value, dtype=np_dtype))
-      elif np_dtype == np.dtype(object) or np.issubdtype(np_dtype, np.bytes_):
-        arrays.append(np.array(feature.bytes_list.value, dtype=np_dtype))
+      if dtype.is_integer:
+        tensors.append(tf.constant(feature.int64_list.value, dtype=dtype))
+      elif dtype == tf.string:
+        tensors.append(tf.constant(feature.bytes_list.value, dtype=dtype))
       else:
-        raise RuntimeError("Unexpected numpy dtype: '{}'.".format(np_dtype))
-    return arrays
+        raise RuntimeError(f"Unexpected dtype: '{dtype}'.")
+    return tensors

tensorflow_compression/python/layers/signal_conv.py

@@ -14,33 +14,25 @@
 # ==============================================================================
 """Tests of PackedTensors class."""
 
-import numpy as np
 import tensorflow as tf
 from tensorflow_compression.python.util import packed_tensors
 
 
 class PackedTensorsTest(tf.test.TestCase):
 
-  def test_pack_unpack(self):
-    """Tests packing and unpacking tensors."""
-    string = np.array(["xyz".encode("ascii")], dtype=object)
-    shape = np.array([1, 3], dtype=np.int32)
-    arrays = [string, shape]
-
-    string_t = tf.zeros([1], dtype=tf.string)
-    shape_t = tf.zeros([2], dtype=tf.int32)
-    tensors = [string_t, shape_t]
-
+  def test_pack_unpack_identity(self):
+    """Tests packing and unpacking tensors returns the same values."""
+    string = tf.constant(["xyz"], dtype=tf.string)
+    shape = tf.constant([1, 3], dtype=tf.int32)
     packed = packed_tensors.PackedTensors()
-    packed.pack(tensors, arrays)
+    packed.pack([string, shape])
     packed = packed_tensors.PackedTensors(packed.string)
-    string_u, shape_u = packed.unpack(tensors)
-
-    self.assertAllEqual(string_u, string)
-    self.assertAllEqual(shape_u, shape)
+    string_unpacked, shape_unpacked = packed.unpack([tf.string, tf.int32])
+    self.assertAllEqual(string_unpacked, string)
+    self.assertAllEqual(shape_unpacked, shape)
 
-  def test_model(self):
-    """Tests setting and getting model."""
+  def test_set_get_model_identity(self):
+    """Tests setting and getting model returns the same value."""
     packed = packed_tensors.PackedTensors()
     packed.model = "xyz"
     packed = packed_tensors.PackedTensors(packed.string)