Converted unit tests to import top-level modules.

Author: Johannes Ballé

tensorflow_compression/python/layers/entropy_models_test.py

@@ -22,26 +22,22 @@
 # Dependency imports
 
 import numpy as np
+import tensorflow as tf
 
-from tensorflow.python.framework import dtypes
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
-from tensorflow.python.training import gradient_descent
 
-from tensorflow_compression.python.layers import entropy_models
+import tensorflow_compression as tfc
 
 
 class EntropyBottleneckTest(test.TestCase):
 
   def test_noise(self):
     # Tests that the noise added is uniform noise between -0.5 and 0.5.
-    inputs = array_ops.placeholder(dtypes.float32, (None, 1))
-    layer = entropy_models.EntropyBottleneck()
+    inputs = tf.placeholder(tf.float32, (None, 1))
+    layer = tfc.EntropyBottleneck()
     noisy, _ = layer(inputs, training=True)
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       values = np.linspace(-50, 50, 100)[:, None]
       noisy, = sess.run([noisy], {inputs: values})
       self.assertFalse(np.allclose(values, noisy, rtol=0, atol=.49))
@@ -50,14 +46,14 @@ def test_noise(self):
   def test_quantization(self):
     # Tests that inputs are quantized to full integer values, even after
     # quantiles have been updated.
-    inputs = array_ops.placeholder(dtypes.float32, (None, 1))
-    layer = entropy_models.EntropyBottleneck(optimize_integer_offset=False)
+    inputs = tf.placeholder(tf.float32, (None, 1))
+    layer = tfc.EntropyBottleneck(optimize_integer_offset=False)
     quantized, _ = layer(inputs, training=False)
-    opt = gradient_descent.GradientDescentOptimizer(learning_rate=1)
+    opt = tf.train.GradientDescentOptimizer(learning_rate=1)
     self.assertTrue(len(layer.losses) == 1)
     step = opt.minimize(layer.losses[0])
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       sess.run(step)
       values = np.linspace(-50, 50, 100)[:, None]
       quantized, = sess.run([quantized], {inputs: values})
@@ -67,14 +63,14 @@ def test_quantization_optimized_offset(self):
     # Tests that inputs are not quantized to full integer values after quantiles
     # have been updated. However, the difference between input and output should
     # be between -0.5 and 0.5, and the offset must be consistent.
-    inputs = array_ops.placeholder(dtypes.float32, (None, 1))
-    layer = entropy_models.EntropyBottleneck(optimize_integer_offset=True)
+    inputs = tf.placeholder(tf.float32, (None, 1))
+    layer = tfc.EntropyBottleneck(optimize_integer_offset=True)
     quantized, _ = layer(inputs, training=False)
-    opt = gradient_descent.GradientDescentOptimizer(learning_rate=1)
+    opt = tf.train.GradientDescentOptimizer(learning_rate=1)
     self.assertTrue(len(layer.losses) == 1)
     step = opt.minimize(layer.losses[0])
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       sess.run(step)
       values = np.linspace(-50, 50, 100)[:, None]
       quantized, = sess.run([quantized], {inputs: values})
@@ -86,17 +82,17 @@ def test_quantization_optimized_offset(self):
   def test_codec(self):
     # Tests that inputs are compressed and decompressed correctly, and quantized
     # to full integer values, even after quantiles have been updated.
-    inputs = array_ops.placeholder(dtypes.float32, (1, None, 1))
-    layer = entropy_models.EntropyBottleneck(
+    inputs = tf.placeholder(tf.float32, (1, None, 1))
+    layer = tfc.EntropyBottleneck(
         data_format="channels_last", init_scale=60,
         optimize_integer_offset=False)
     bitstrings = layer.compress(inputs)
-    decoded = layer.decompress(bitstrings, array_ops.shape(inputs)[1:])
-    opt = gradient_descent.GradientDescentOptimizer(learning_rate=1)
+    decoded = layer.decompress(bitstrings, tf.shape(inputs)[1:])
+    opt = tf.train.GradientDescentOptimizer(learning_rate=1)
     self.assertTrue(len(layer.losses) == 1)
     step = opt.minimize(layer.losses[0])
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       sess.run(step)
       self.assertTrue(len(layer.updates) == 1)
       sess.run(layer.updates[0])
@@ -109,17 +105,17 @@ def test_codec_optimized_offset(self):
     # quantized to full integer values after quantiles have been updated.
     # However, the difference between input and output should be between -0.5
     # and 0.5, and the offset must be consistent.
-    inputs = array_ops.placeholder(dtypes.float32, (1, None, 1))
-    layer = entropy_models.EntropyBottleneck(
+    inputs = tf.placeholder(tf.float32, (1, None, 1))
+    layer = tfc.EntropyBottleneck(
         data_format="channels_last", init_scale=60,
         optimize_integer_offset=True)
     bitstrings = layer.compress(inputs)
-    decoded = layer.decompress(bitstrings, array_ops.shape(inputs)[1:])
-    opt = gradient_descent.GradientDescentOptimizer(learning_rate=1)
+    decoded = layer.decompress(bitstrings, tf.shape(inputs)[1:])
+    opt = tf.train.GradientDescentOptimizer(learning_rate=1)
     self.assertTrue(len(layer.losses) == 1)
     step = opt.minimize(layer.losses[0])
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       sess.run(step)
       self.assertTrue(len(layer.updates) == 1)
       sess.run(layer.updates[0])
@@ -133,13 +129,13 @@ def test_codec_optimized_offset(self):
   def test_codec_clipping(self):
     # Tests that inputs are compressed and decompressed correctly, and clipped
     # to the expected range.
-    inputs = array_ops.placeholder(dtypes.float32, (1, None, 1))
-    layer = entropy_models.EntropyBottleneck(
+    inputs = tf.placeholder(tf.float32, (1, None, 1))
+    layer = tfc.EntropyBottleneck(
         data_format="channels_last", init_scale=40)
     bitstrings = layer.compress(inputs)
-    decoded = layer.decompress(bitstrings, array_ops.shape(inputs)[1:])
+    decoded = layer.decompress(bitstrings, tf.shape(inputs)[1:])
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       self.assertTrue(len(layer.updates) == 1)
       sess.run(layer.updates[0])
       values = np.linspace(-50, 50, 100)[None, :, None]
@@ -150,15 +146,15 @@ def test_codec_clipping(self):
   def test_channels_last(self):
     # Test the layer with more than one channel and multiple input dimensions,
     # with the channels in the last dimension.
-    inputs = array_ops.placeholder(dtypes.float32, (None, None, None, 2))
-    layer = entropy_models.EntropyBottleneck(
+    inputs = tf.placeholder(tf.float32, (None, None, None, 2))
+    layer = tfc.EntropyBottleneck(
         data_format="channels_last", init_scale=50)
     noisy, _ = layer(inputs, training=True)
     quantized, _ = layer(inputs, training=False)
     bitstrings = layer.compress(inputs)
-    decoded = layer.decompress(bitstrings, array_ops.shape(inputs)[1:])
+    decoded = layer.decompress(bitstrings, tf.shape(inputs)[1:])
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       self.assertTrue(len(layer.updates) == 1)
       sess.run(layer.updates[0])
       values = 5 * np.random.normal(size=(7, 5, 3, 2))
@@ -171,15 +167,15 @@ def test_channels_last(self):
   def test_channels_first(self):
     # Test the layer with more than one channel and multiple input dimensions,
     # with the channel dimension right after the batch dimension.
-    inputs = array_ops.placeholder(dtypes.float32, (None, 3, None, None))
-    layer = entropy_models.EntropyBottleneck(
+    inputs = tf.placeholder(tf.float32, (None, 3, None, None))
+    layer = tfc.EntropyBottleneck(
         data_format="channels_first", init_scale=50)
     noisy, _ = layer(inputs, training=True)
     quantized, _ = layer(inputs, training=False)
     bitstrings = layer.compress(inputs)
-    decoded = layer.decompress(bitstrings, array_ops.shape(inputs)[1:])
+    decoded = layer.decompress(bitstrings, tf.shape(inputs)[1:])
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       self.assertTrue(len(layer.updates) == 1)
       sess.run(layer.updates[0])
       values = 5 * np.random.normal(size=(2, 3, 5, 7))
@@ -193,13 +189,13 @@ def test_compress(self):
     # Test compression and decompression, and produce test data for
     # `test_decompress`. If you set the constant at the end to `True`, this test
     # will fail and the log will contain the new test data.
-    inputs = array_ops.placeholder(dtypes.float32, (2, 3, 10))
-    layer = entropy_models.EntropyBottleneck(
+    inputs = tf.placeholder(tf.float32, (2, 3, 10))
+    layer = tfc.EntropyBottleneck(
         data_format="channels_first", filters=(), init_scale=2)
     bitstrings = layer.compress(inputs)
-    decoded = layer.decompress(bitstrings, array_ops.shape(inputs)[1:])
+    decoded = layer.decompress(bitstrings, tf.shape(inputs)[1:])
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       self.assertTrue(len(layer.updates) == 1)
       sess.run(layer.updates[0])
       values = 5 * np.random.uniform(size=(2, 3, 10)) - 2.5
@@ -236,54 +232,54 @@ def test_compress(self):
   def test_decompress(self):
     # Test that decompression of values compressed with a previous version
     # works, i.e. that the file format doesn't change across revisions.
-    bitstrings = array_ops.placeholder(dtypes.string)
-    input_shape = array_ops.placeholder(dtypes.int32)
-    quantized_cdf = array_ops.placeholder(dtypes.int32)
-    layer = entropy_models.EntropyBottleneck(
-        data_format="channels_first", filters=(), dtype=dtypes.float32)
+    bitstrings = tf.placeholder(tf.string)
+    input_shape = tf.placeholder(tf.int32)
+    quantized_cdf = tf.placeholder(tf.int32)
+    layer = tfc.EntropyBottleneck(
+        data_format="channels_first", filters=(), dtype=tf.float32)
     layer.build(self.expected.shape)
     layer._quantized_cdf = quantized_cdf
     decoded = layer.decompress(bitstrings, input_shape[1:])
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       decoded, = sess.run([decoded], {
           bitstrings: self.bitstrings, input_shape: self.expected.shape,
           quantized_cdf: self.quantized_cdf})
       self.assertAllClose(self.expected, decoded, rtol=0, atol=1e-6)
 
   def test_build_decompress(self):
     # Test that layer can be built when `decompress` is the first call to it.
-    bitstrings = array_ops.placeholder(dtypes.string)
-    input_shape = array_ops.placeholder(dtypes.int32, shape=[3])
-    layer = entropy_models.EntropyBottleneck(dtype=dtypes.float32)
+    bitstrings = tf.placeholder(tf.string)
+    input_shape = tf.placeholder(tf.int32, shape=[3])
+    layer = tfc.EntropyBottleneck(dtype=tf.float32)
     layer.decompress(bitstrings, input_shape[1:], channels=5)
     self.assertTrue(layer.built)
 
   def test_pmf_normalization(self):
     # Test that probability mass functions are normalized correctly.
-    layer = entropy_models.EntropyBottleneck(dtype=dtypes.float32)
+    layer = tfc.EntropyBottleneck(dtype=tf.float32)
     layer.build((None, 10))
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       pmf, = sess.run([layer._pmf])
       self.assertAllClose(np.ones(10), np.sum(pmf, axis=-1), rtol=0, atol=1e-6)
 
   def test_visualize(self):
     # Test that summary op can be constructed.
-    layer = entropy_models.EntropyBottleneck(dtype=dtypes.float32)
+    layer = tfc.EntropyBottleneck(dtype=tf.float32)
     layer.build((None, 10))
     summary = layer.visualize()
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       sess.run([summary])
 
   def test_normalization(self):
     # Test that densities are normalized correctly.
-    inputs = array_ops.placeholder(dtypes.float32, (None, 1))
-    layer = entropy_models.EntropyBottleneck(filters=(2,))
+    inputs = tf.placeholder(tf.float32, (None, 1))
+    layer = tfc.EntropyBottleneck(filters=(2,))
     _, likelihood = layer(inputs, training=True)
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       x = np.repeat(np.arange(-200, 201), 1000)[:, None]
       likelihood, = sess.run([likelihood], {inputs: x})
       self.assertEqual(x.shape, likelihood.shape)
@@ -292,16 +288,16 @@ def test_normalization(self):
 
   def test_entropy_estimates(self):
     # Test that entropy estimates match actual range coding.
-    inputs = array_ops.placeholder(dtypes.float32, (1, None, 1))
-    layer = entropy_models.EntropyBottleneck(
+    inputs = tf.placeholder(tf.float32, (1, None, 1))
+    layer = tfc.EntropyBottleneck(
         filters=(2, 3), data_format="channels_last")
     _, likelihood = layer(inputs, training=True)
-    diff_entropy = math_ops.reduce_sum(math_ops.log(likelihood)) / -np.log(2)
+    diff_entropy = tf.reduce_sum(tf.log(likelihood)) / -np.log(2)
     _, likelihood = layer(inputs, training=False)
-    disc_entropy = math_ops.reduce_sum(math_ops.log(likelihood)) / -np.log(2)
+    disc_entropy = tf.reduce_sum(tf.log(likelihood)) / -np.log(2)
     bitstrings = layer.compress(inputs)
     with self.test_session() as sess:
-      sess.run(variables.global_variables_initializer())
+      sess.run(tf.global_variables_initializer())
       self.assertTrue(len(layer.updates) == 1)
       sess.run(layer.updates[0])
       diff_entropy, disc_entropy, bitstrings = sess.run(

tensorflow_compression/python/layers/gdn_test.py

@@ -23,15 +23,14 @@
 
 import numpy as np
 import tensorflow as tf
-
-from tensorflow_compression.python.layers import gdn
+import tensorflow_compression as tfc
 
 
 class GDNTest(tf.test.TestCase):
 
   def _run_gdn(self, x, shape, inverse, rectify, data_format):
     inputs = tf.placeholder(tf.float32, shape)
-    layer = gdn.GDN(
+    layer = tfc.GDN(
         inverse=inverse, rectify=rectify, data_format=data_format)
     outputs = layer(inputs)
     with self.test_session() as sess:

tensorflow_compression/python/layers/parameterizers_test.py

@@ -23,8 +23,7 @@
 
 import numpy as np
 import tensorflow as tf
-
-from tensorflow_compression.python.layers import parameterizers
+import tensorflow_compression as tfc
 
 
 class ParameterizersTest(tf.test.TestCase):
@@ -41,31 +40,31 @@ def _test_parameterizer(self, param, init, shape):
   def test_static_parameterizer(self):
     shape = (1, 2, 3, 4)
     var = self._test_parameterizer(
-        parameterizers.StaticParameterizer(tf.initializers.zeros()),
+        tfc.StaticParameterizer(tf.initializers.zeros()),
         tf.initializers.random_uniform(), shape)
     self.assertEqual(var.shape, shape)
     self.assertAllClose(var, np.zeros(shape), rtol=0, atol=1e-7)
 
   def test_rdft_parameterizer(self):
     shape = (3, 4, 2, 1)
     var = self._test_parameterizer(
-        parameterizers.RDFTParameterizer(),
+        tfc.RDFTParameterizer(),
         tf.initializers.ones(), shape)
     self.assertEqual(var.shape, shape)
     self.assertAllClose(var, np.ones(shape), rtol=0, atol=1e-6)
 
   def test_nonnegative_parameterizer(self):
     shape = (1, 2, 3, 4)
     var = self._test_parameterizer(
-        parameterizers.NonnegativeParameterizer(),
+        tfc.NonnegativeParameterizer(),
         tf.initializers.random_uniform(), shape)
     self.assertEqual(var.shape, shape)
     self.assertTrue(np.all(var >= 0))
 
   def test_positive_parameterizer(self):
     shape = (1, 2, 3, 4)
     var = self._test_parameterizer(
-        parameterizers.NonnegativeParameterizer(minimum=.1),
+        tfc.NonnegativeParameterizer(minimum=.1),
         tf.initializers.random_uniform(), shape)
     self.assertEqual(var.shape, shape)
     self.assertTrue(np.all(var >= .1))

tensorflow_compression/python/layers/signal_test.py

@@ -26,9 +26,8 @@
 import tensorflow as tf
 
 from tensorflow.python.platform import test
-from tensorflow_compression.python.layers import initializers
-from tensorflow_compression.python.layers import parameterizers
-from tensorflow_compression.python.layers import signal
+
+import tensorflow_compression as tfc
 
 
 class SignalTest(tf.test.TestCase):
@@ -98,14 +97,13 @@ def run_valid(self, batch, input_support, channels, filters, kernel_support,
     # Create kernel array.
     kernel = np.random.randint(16, size=kernel_support + (channels, filters))
     kernel = kernel.astype(np.float32)
-    tf_kernel = parameterizers.StaticParameterizer(
-        tf.constant_initializer(kernel))
+    tf_kernel = tfc.StaticParameterizer(tf.constant_initializer(kernel))
 
     # Run SignalConv* layer.
     layer_class = {
-        3: signal.SignalConv1D,
-        4: signal.SignalConv2D,
-        5: signal.SignalConv3D,
+        3: tfc.SignalConv1D,
+        4: tfc.SignalConv2D,
+        5: tfc.SignalConv3D,
     }[inputs.ndim]
     layer = layer_class(
         filters, kernel_support, corr=corr, strides_down=strides_down,
@@ -148,14 +146,13 @@ def run_same(self, batch, input_support, channels, filters, kernel_support,
 
     # Create kernel array. This is an identity kernel, so the outputs should
     # be equal to the inputs except for up- and downsampling.
-    tf_kernel = parameterizers.StaticParameterizer(
-        initializers.IdentityInitializer())
+    tf_kernel = tfc.StaticParameterizer(tfc.IdentityInitializer())
 
     # Run SignalConv* layer.
     layer_class = {
-        3: signal.SignalConv1D,
-        4: signal.SignalConv2D,
-        5: signal.SignalConv3D,
+        3: tfc.SignalConv1D,
+        4: tfc.SignalConv2D,
+        5: tfc.SignalConv3D,
     }[inputs.ndim]
     layer = layer_class(
         1, kernel_support, corr=corr, strides_down=strides_down,