Reimplements RDFT kernel reparameterization using tf.signal.

PiperOrigin-RevId: 367329695
Change-Id: Ibcbc6cc96b48560494a3b9e2c0e2305c90e7a138

Author: Johannes Ballé

BUILD

@@ -25,7 +25,6 @@ py_library(
         "//tensorflow_compression/python/ops:math_ops",
         "//tensorflow_compression/python/ops:padding_ops",
         "//tensorflow_compression/python/ops:soft_round_ops",
-        "//tensorflow_compression/python/ops:spectral_ops",
         "//tensorflow_compression/python/util:packed_tensors",
     ],
 )

tensorflow_compression/__init__.py

@@ -36,7 +36,6 @@
 from tensorflow_compression.python.ops.math_ops import *
 from tensorflow_compression.python.ops.padding_ops import *
 from tensorflow_compression.python.ops.soft_round_ops import *
-from tensorflow_compression.python.ops.spectral_ops import *
 
 from tensorflow_compression.python.util.packed_tensors import *
 # pylint: enable=wildcard-import

tensorflow_compression/all_tests.py

@@ -41,7 +41,6 @@
 from tensorflow_compression.python.ops.padding_ops_test import *
 from tensorflow_compression.python.ops.range_coding_ops_test import *
 from tensorflow_compression.python.ops.soft_round_ops_test import *
-from tensorflow_compression.python.ops.spectral_ops_test import *
 
 from tensorflow_compression.python.util.packed_tensors_test import *
 # pylint: enable=wildcard-import

tensorflow_compression/python/layers/BUILD

@@ -38,10 +38,7 @@ py_library(
     name = "parameters",
     srcs = ["parameters.py"],
     srcs_version = "PY3",
-    deps = [
-        "//tensorflow_compression/python/ops:math_ops",
-        "//tensorflow_compression/python/ops:spectral_ops",
-    ],
+    deps = ["//tensorflow_compression/python/ops:math_ops"],
 )
 
 py_test(

tensorflow_compression/python/layers/parameters.py

@@ -18,7 +18,6 @@
 from typing import Any, Dict
 import tensorflow as tf
 from tensorflow_compression.python.ops import math_ops
-from tensorflow_compression.python.ops import spectral_ops
 
 
 __all__ = [
@@ -78,50 +77,64 @@ class RDFTParameter(Parameter):
   (see https://en.wikipedia.org/wiki/Discrete_Fourier_transform)
 
   Attributes:
-    dc: Boolean. The `dc` parameter provided on initialization.
     shape: `tf.TensorShape`. The shape of the convolution kernel.
-    rdft: `tf.Variable`. The RDFT of the kernel.
+    real: `tf.Variable`. The real part of the RDFT of the kernel.
+    imag: `tf.Variable`. The imaginary part of the RDFT of the kernel.
   """
 
-  def __init__(self, initial_value, name=None, dc=True, shape=None, dtype=None):
+  def __init__(self, initial_value, name=None, shape=None, dtype=None):
     """Initializer.
 
     Args:
       initial_value: `tf.Tensor` or `None`. The initial value of the kernel. If
         not provided, its `shape` must be given, and the initial value of the
         parameter will be undefined.
       name: String. The name of the kernel.
-      dc: Boolean. If `False`, the DC component of the kernel RDFTs is not
-        represented, forcing the filters to be highpass. Defaults to `True`.
       shape: `tf.TensorShape` or compatible. Ignored unless `initial_value is
         None`.
       dtype: `tf.dtypes.DType` or compatible. DType of this parameter. If not
         given, inferred from `initial_value`.
     """
     super().__init__(name=name)
-    self._dc = bool(dc)
     if initial_value is None:
       if shape is None:
         raise ValueError("If initial_value is None, shape must be specified.")
       initial_value = tf.zeros(shape, dtype=dtype)
     else:
       initial_value = tf.convert_to_tensor(initial_value, dtype=dtype)
     self._shape = initial_value.shape
-    self._matrix = spectral_ops.irdft_matrix(
-        self.shape[:-2], dtype=initial_value.dtype)
-    if not self.dc:
-      self._matrix = self._matrix[:, 1:]
-    initial_value = tf.reshape(
-        initial_value, (-1, self.shape[-2] * self.shape[-1]))
-    initial_value = tf.linalg.matmul(
-        self._matrix, initial_value, transpose_a=True)
+    self._dtype = initial_value.dtype
+
+    if self.shape.rank == 3:
+      initial_value = tf.transpose(initial_value, (1, 2, 0))
+      initial_value = tf.signal.rfft(initial_value)
+    elif self.shape.rank == 4:
+      initial_value = tf.transpose(initial_value, (2, 3, 0, 1))
+      initial_value = tf.signal.rfft2d(initial_value)
+    elif self.shape.rank == 5:
+      initial_value = tf.transpose(initial_value, (3, 4, 0, 1, 2))
+      initial_value = tf.signal.rfft3d(initial_value)
+    else:
+      raise ValueError(
+          f"Expected kernel tensor of rank 3, 4, or 5; received shape "
+          f"{self._shape}.")
+    self._norm = tf.constant(
+        self.shape[:-2].num_elements() ** .5, initial_value.dtype)
+    initial_value /= self._norm
+    # We split the variable into real and imaginary parts to avoid issues with
+    # complex-valued variables being unsupported when saving models, etc.
+    real = tf.math.real(initial_value)
+    imag = tf.math.imag(initial_value)
+    real_name = imag_name = None
     if name is not None:
-      name = f"{name}_rdft"
-    self.rdft = tf.Variable(initial_value, name=name)
+      real_name = f"{name}_real"
+      imag_name = f"{name}_imag"
+    self.real = tf.Variable(real, name=real_name)
+    self.imag = tf.Variable(imag, name=imag_name)
 
   @property
-  def dc(self) -> bool:
-    return self._dc
+  def dtype(self) -> tf.dtypes.DType:
+    return self._dtype
 
   @property
   def shape(self) -> tf.TensorShape:
@@ -130,15 +143,24 @@ def shape(self) -> tf.TensorShape:
   @tf.Module.with_name_scope
   def __call__(self) -> tf.Tensor:
     """Computes and returns the convolution kernel as a `tf.Tensor`."""
-    return tf.reshape(tf.linalg.matmul(self._matrix, self.rdft), self.shape)
+    rdft = tf.dtypes.complex(self.real, self.imag) * self._norm
+    if self.shape.rank == 3:
+      kernel = tf.signal.irfft(rdft, fft_length=self.shape[:-2])
+      return tf.transpose(kernel, (2, 0, 1))
+    elif self.shape.rank == 4:
+      kernel = tf.signal.irfft2d(rdft, fft_length=self.shape[:-2])
+      return tf.transpose(kernel, (2, 3, 0, 1))
+    else:
+      assert self.shape.rank == 5, self.shape
+      kernel = tf.signal.irfft3d(rdft, fft_length=self.shape[:-2])
+      return tf.transpose(kernel, (2, 3, 4, 0, 1))
 
   def get_config(self) -> Dict[str, Any]:
     config = super().get_config()
     config.update(
         initial_value=None,
-        dc=self.dc,
-        shape=tuple(self.shape),
-        dtype=self.rdft.dtype.name,
+        shape=tuple(map(int, self.shape)),
+        dtype=self.dtype.name,
     )
     return config
 
@@ -220,7 +242,7 @@ def get_config(self) -> Dict[str, Any]:
         initial_value=None,
         minimum=self.minimum,
         offset=self.offset,
-        shape=tuple(self.variable.shape),
+        shape=tuple(map(int, self.variable.shape)),
         dtype=self.variable.dtype.name,
     )
     return config

tensorflow_compression/python/layers/parameters_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Tests of parameters."""
 
+from absl.testing import parameterized
 import tensorflow as tf
 from tensorflow_compression.python.layers import parameters
 
@@ -49,24 +50,33 @@ def test_converts_to_tensor(self):
     self.assertEqual(value.dtype.name, converted.dtype.name)
 
 
-class RDFTParameterTest(ParameterTest, tf.test.TestCase):
+class RDFTParameterTest(ParameterTest, tf.test.TestCase,
+                        parameterized.TestCase):
 
   cls = parameters.RDFTParameter
-  kwargs = dict(name="hello_rdft", dc=True)
+  kwargs = dict(name="rdft_kernel")
   shape = (3, 3, 1, 2)
 
-  def test_initial_value_is_reproduced_without_dc(self):
-    initial_value = tf.random.uniform(self.shape, dtype=tf.float32)
-    parameter = self.cls(initial_value, dc=False)
-    expected_value = initial_value - tf.reduce_mean(
-        initial_value, axis=(0, 1), keepdims=True)
-    self.assertAllClose(expected_value, parameter(), atol=1e-6, rtol=0)
+  # TODO(jonycgn): Find out why 3D RFFT gradients are not implemented in TF.
+  @parameterized.parameters((7, 3, 2), (5, 3, 1, 2))
+  def test_gradients_propagate(self, *shape):
+    initial_value = tf.random.uniform(shape, dtype=tf.float32)
+    parameter = self.cls(initial_value, **self.kwargs)
+    rand = tf.random.uniform(shape)
+    with tf.GradientTape() as tape:
+      loss = tf.reduce_sum(rand * parameter())
+    gradients = tape.gradient(loss, parameter.variables)
+    self.assertLen(gradients, 2)
+    self.assertNotAllClose(
+        tf.zeros_like(gradients[0]), gradients[0], atol=1e-1, rtol=0)
+    self.assertNotAllClose(
+        tf.zeros_like(gradients[1]), gradients[1], atol=1e-1, rtol=0)
 
 
 class GDNParameterTest(ParameterTest, tf.test.TestCase):
 
   cls = parameters.GDNParameter
-  kwargs = dict(name="hello_gdn")
+  kwargs = dict(name="gdn_parameter")
   shape = (2, 1, 3)
 
   def test_initial_value_is_reproduced_with_minimum(self):

tensorflow_compression/python/layers/signal_conv_test.py

@@ -33,18 +33,19 @@ def test_invalid_data_format_raises_error(self):
   def test_variables_are_enumerated(self):
     layer = signal_conv.SignalConv2D(3, 1, use_bias=True)
     layer.build((None, None, None, 2))
-    self.assertLen(layer.weights, 2)
-    self.assertLen(layer.trainable_weights, 2)
+    self.assertLen(layer.weights, 3)
+    self.assertLen(layer.trainable_weights, 3)
     weight_names = [w.name for w in layer.weights]
-    self.assertSameElements(weight_names, ["kernel_rdft:0", "bias:0"])
+    self.assertSameElements(
+        weight_names, ["kernel_real:0", "kernel_imag:0", "bias:0"])
 
   def test_bias_variable_is_not_unnecessarily_created(self):
     layer = signal_conv.SignalConv2D(5, 3, use_bias=False)
     layer.build((None, None, None, 3))
-    self.assertLen(layer.weights, 1)
-    self.assertLen(layer.trainable_weights, 1)
+    self.assertLen(layer.weights, 2)
+    self.assertLen(layer.trainable_weights, 2)
     weight_names = [w.name for w in layer.weights]
-    self.assertSameElements(weight_names, ["kernel_rdft:0"])
+    self.assertSameElements(weight_names, ["kernel_real:0", "kernel_imag:0"])
 
   def test_variables_are_not_enumerated_when_overridden(self):
     layer = signal_conv.SignalConv2D(1, 1)
@@ -58,7 +59,7 @@ def test_variables_trainable_state_follows_layer(self):
     layer = signal_conv.SignalConv2D(1, 1, use_bias=True)
     layer.trainable = False
     layer.build((None, None, None, 1))
-    self.assertLen(layer.weights, 2)
+    self.assertLen(layer.weights, 3)
     self.assertEmpty(layer.trainable_weights)
 
   def test_attributes_cannot_be_set_after_build(self):
@@ -107,7 +108,7 @@ def test_variables_receive_gradients(self):
     with tf.GradientTape() as g:
       y = layer(x)
     grads = g.gradient(y, layer.trainable_weights)
-    self.assertLen(grads, 2)
+    self.assertLen(grads, 3)
     self.assertNotIn(None, grads)
     grad_shapes = [tuple(g.shape) for g in grads]
     weight_shapes = [tuple(w.shape) for w in layer.trainable_weights]

tensorflow_compression/python/ops/BUILD

@@ -60,19 +60,6 @@ py_test(
     deps = [":soft_round_ops"],
 )
 
-py_library(
-    name = "spectral_ops",
-    srcs = ["spectral_ops.py"],
-    srcs_version = "PY3",
-)
-
-py_test(
-    name = "spectral_ops_test",
-    srcs = ["spectral_ops_test.py"],
-    python_version = "PY3",
-    deps = [":spectral_ops"],
-)
-
 filegroup(
     name = "py_src",
     srcs = glob(["*.py"]),