Updates initializers to TF2/Keras API and adds unit tests.

PiperOrigin-RevId: 354464149
Change-Id: Ibd3fdfbc0e581c15b1bd9c4404e9b159330b6a20

Author: Johannes Ballé

tensorflow_compression/all_tests.py

@@ -31,6 +31,7 @@
 
 from tensorflow_compression.python.layers.entropy_models_test import *
 from tensorflow_compression.python.layers.gdn_test import *
+from tensorflow_compression.python.layers.initializers_test import *
 from tensorflow_compression.python.layers.parameterizers_test import *
 from tensorflow_compression.python.layers.signal_conv_test import *
 from tensorflow_compression.python.layers.soft_round_test import *

tensorflow_compression/python/layers/BUILD

@@ -68,6 +68,13 @@ py_test(
     deps = [":entropy_models"],
 )
 
+py_test(
+    name = "initializers_test",
+    srcs = ["initializers_test.py"],
+    python_version = "PY3",
+    deps = [":initializers"],
+)
+
 py_test(
     name = "parameterizers_test",
     srcs = ["parameterizers_test.py"],

tensorflow_compression/python/layers/initializers.py

@@ -14,39 +14,42 @@
 # ==============================================================================
 """Initializers for layer classes."""
 
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 
 
 __all__ = [
     "IdentityInitializer",
 ]
 
 
-class IdentityInitializer(object):
+class IdentityInitializer(tf.keras.initializers.Initializer):
   """Initialize to the identity kernel with the given shape.
 
   This creates an n-D kernel suitable for `SignalConv*` with the requested
   support that produces an output identical to its input (except possibly at the
   signal boundaries).
 
-  Note: The identity initializer in `tf.initializers` is only suitable for
+  Note: The identity initializer in `tf.keras.initializers` is only suitable for
   matrices, not for n-D convolution kernels (i.e., no spatial support).
   """
 
   def __init__(self, gain=1):
-    self.gain = float(gain)
+    super().__init__()
+    self.gain = gain
 
-  def __call__(self, shape, dtype=None, partition_info=None):
-    del partition_info  # unused
-    assert len(shape) > 2, shape
+  def __call__(self, shape, dtype=None, **kwargs):
+    del kwargs  # unused
+    shape = tf.TensorShape(shape)
+    if shape.rank <= 2:
+      raise ValueError(f"shape must be at least rank 3, got {shape}.")
 
-    support = tuple(shape[:-2]) + (1, 1)
+    support = shape.as_list()[:-2] + [1, 1]
     indices = [[s // 2 for s in support]]
     updates = tf.constant([self.gain], dtype=dtype)
-    kernel = tf.scatter_nd(indices, updates, support)
+    spatial_kernel = tf.scatter_nd(indices, updates, support)
+    return spatial_kernel * tf.eye(shape[-2], shape[-1], dtype=dtype)
 
-    assert shape[-2] == shape[-1], shape
-    if shape[-1] != 1:
-      kernel *= tf.eye(shape[-1], dtype=dtype)
-
-    return kernel
+  def get_config(self):
+    config = super().get_config()
+    config.update(gain=self.gain)
+    return config

tensorflow_compression/python/layers/initializers_test.py

@@ -0,0 +1,49 @@
+# Copyright 2021 Google LLC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests of initializers."""
+
+import tensorflow as tf
+from tensorflow_compression.python.layers import initializers
+
+
+class InitializerTest(tf.test.TestCase):
+
+  def test_creates_1d_kernel(self):
+    expected_kernel = tf.transpose([
+        [[0, 3, 0], [0, 0, 0], [0, 0, 0]],
+        [[0, 0, 0], [0, 3, 0], [0, 0, 0]],
+        [[0, 0, 0], [0, 0, 0], [0, 3, 0]],
+        [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
+    ], (2, 0, 1))
+    kernel = initializers.IdentityInitializer(gain=3)((3, 4, 3), dtype=tf.int32)
+    self.assertAllEqual(expected_kernel, kernel)
+
+  def test_creates_2d_kernel(self):
+    expected_kernel = tf.constant([
+        [0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0],
+        [0, 0, 1, 0, 0],
+        [0, 0, 0, 0, 0],
+    ])[:, :, None, None]
+    kernel = initializers.IdentityInitializer()((4, 5, 1, 1), dtype=tf.float32)
+    self.assertAllEqual(expected_kernel, kernel)
+
+  def test_fails_for_invalid_shape(self):
+    with self.assertRaises(ValueError):
+      initializers.IdentityInitializer()((2, 3), dtype=tf.float32)
+
+
+if __name__ == "__main__":
+  tf.test.main()