Fixes deprecated names.

- In Keras layers, replaces add_variable() with add_weight().
- In Keras layers, removes inputs= argument in add_update().
- In unit tests, replaces test_session() with cached_session().

PiperOrigin-RevId: 281393590
Change-Id: I7c66f654d3828a657a5f55c06fdd00f510ef1043

Author: Johannes Ball?

tensorflow_compression/python/layers/entropy_models.py

@@ -467,21 +467,21 @@ def build(self, input_shape):
     self._factors = []
     for i in range(len(self.filters) + 1):
       init = np.log(np.expm1(1 / scale / filters[i + 1]))
-      matrix = self.add_variable(
+      matrix = self.add_weight(
           "matrix_{}".format(i), dtype=self.dtype,
           shape=(channels, filters[i + 1], filters[i]),
           initializer=tf.initializers.constant(init))
       matrix = tf.nn.softplus(matrix)
       self._matrices.append(matrix)
 
-      bias = self.add_variable(
+      bias = self.add_weight(
           "bias_{}".format(i), dtype=self.dtype,
           shape=(channels, filters[i + 1], 1),
           initializer=tf.initializers.random_uniform(-.5, .5))
       self._biases.append(bias)
 
       if i < len(self.filters):
-        factor = self.add_variable(
+        factor = self.add_weight(
             "factor_{}".format(i), dtype=self.dtype,
             shape=(channels, filters[i + 1], 1),
             initializer=tf.initializers.zeros())
@@ -510,7 +510,7 @@ def quantiles_initializer(shape, dtype=None, partition_info=None):
           [[[-self.init_scale, 0, self.init_scale]]], dtype=dtype)
       return tf.tile(init, (shape[0], 1, 1))
 
-    quantiles = self.add_variable(
+    quantiles = self.add_weight(
         "quantiles", shape=(channels, 1, 3), dtype=self.dtype,
         initializer=quantiles_initializer)
     logits = self._logits_cumulative(quantiles, stop_gradient=True)
@@ -572,13 +572,13 @@ def cdf_initializer(shape, dtype=None, partition_info=None):
     # We need to supply an initializer without fully defined static shape
     # here, or the variable will return the wrong dynamic shape later. A
     # placeholder with default gets the trick done (see initializer above).
-    quantized_cdf = self.add_variable(
+    quantized_cdf = self.add_weight(
         "quantized_cdf",
         shape=(channels, None),
         dtype=tf.int32,
         trainable=False,
         initializer=cdf_initializer)
-    cdf_length = self.add_variable(
+    cdf_length = self.add_weight(
         "cdf_length", shape=(channels,), dtype=tf.int32, trainable=False,
         initializer=tf.initializers.constant(3))
     # Works around a weird TF issue with reading variables inside a loop.
@@ -593,7 +593,7 @@ def cdf_initializer(shape, dtype=None, partition_info=None):
         cdf_length,
         pmf_length + 2)
     update_op = tf.group(update_cdf, update_length)
-    self.add_update(update_op, inputs=None)
+    self.add_update(update_op)
 
     super(EntropyBottleneck, self).build(input_shape)
 
@@ -853,10 +853,10 @@ def cdf_initializer(shape, dtype=None, partition_info=None):
           pmf, tail_mass,
           tf.constant(pmf_length, dtype=tf.int32), max_length)
 
-    quantized_cdf = self.add_variable(
+    quantized_cdf = self.add_weight(
         "quantized_cdf", shape=(len(pmf_length), max_length + 2),
         initializer=cdf_initializer, dtype=tf.int32, trainable=False)
-    cdf_length = self.add_variable(
+    cdf_length = self.add_weight(
         "cdf_length", shape=(len(pmf_length),),
         initializer=tf.initializers.constant(pmf_length + 2),
         dtype=tf.int32, trainable=False)

tensorflow_compression/python/layers/entropy_models_test.py

@@ -32,7 +32,7 @@ def test_noise(self):
     inputs = tf.placeholder(tf.float32, (None, 1))
     layer = entropy_models.EntropyBottleneck()
     noisy, _ = layer(inputs, training=True)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       values = np.linspace(-50, 50, 100)[:, None]
       noisy, = sess.run([noisy], {inputs: values})
@@ -45,7 +45,7 @@ def test_quantization_init(self):
     inputs = tf.placeholder(tf.float32, (None, 1))
     layer = entropy_models.EntropyBottleneck()
     quantized, _ = layer(inputs, training=False)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       values = np.linspace(-50, 50, 100)[:, None]
       quantized, = sess.run([quantized], {inputs: values})
@@ -61,7 +61,7 @@ def test_quantization(self):
     opt = tf.train.GradientDescentOptimizer(learning_rate=1)
     self.assertEqual(1, len(layer.losses))
     step = opt.minimize(layer.losses[0])
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       sess.run(step)
       values = np.linspace(-50, 50, 100)[:, None]
@@ -79,7 +79,7 @@ def test_codec_init(self):
         data_format="channels_last", init_scale=30)
     bitstrings = layer.compress(inputs)
     decoded = layer.decompress(bitstrings, tf.shape(inputs)[1:])
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       values = np.linspace(-50, 50, 100)[None, :, None]
       decoded, = sess.run([decoded], {inputs: values})
@@ -98,7 +98,7 @@ def test_codec(self):
     opt = tf.train.GradientDescentOptimizer(learning_rate=1)
     self.assertEqual(1, len(layer.losses))
     step = opt.minimize(layer.losses[0])
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       sess.run(step)
       self.assertEqual(1, len(layer.updates))
@@ -120,7 +120,7 @@ def test_channels_last(self):
     quantized, _ = layer(inputs, training=False)
     bitstrings = layer.compress(inputs)
     decoded = layer.decompress(bitstrings, tf.shape(inputs)[1:])
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       self.assertEqual(1, len(layer.updates))
       sess.run(layer.updates[0])
@@ -141,7 +141,7 @@ def test_channels_first(self):
     quantized, _ = layer(inputs, training=False)
     bitstrings = layer.compress(inputs)
     decoded = layer.decompress(bitstrings, tf.shape(inputs)[1:])
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       self.assertEqual(1, len(layer.updates))
       sess.run(layer.updates[0])
@@ -161,7 +161,7 @@ def test_compress(self):
         data_format="channels_first", filters=(), init_scale=2)
     bitstrings = layer.compress(inputs)
     decoded = layer.decompress(bitstrings, tf.shape(inputs)[1:])
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       values = 8 * np.random.uniform(size=(2, 3, 9)) - 4
       sess.run(tf.global_variables_initializer())
       self.assertEqual(1, len(layer.updates))
@@ -213,7 +213,7 @@ def test_decompress(self):
     layer._quantized_cdf = quantized_cdf
     layer._cdf_length = cdf_length
     decoded = layer.decompress(bitstrings, input_shape[1:])
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       decoded, = sess.run([decoded], {
           bitstrings: self.bitstrings, input_shape: self.expected.shape,
@@ -233,7 +233,7 @@ def test_normalization(self):
     inputs = tf.placeholder(tf.float32, (None, 1))
     layer = entropy_models.EntropyBottleneck(filters=(2,))
     _, likelihood = layer(inputs, training=True)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       x = np.repeat(np.arange(-200, 201), 2000)[:, None]
       likelihood, = sess.run([likelihood], {inputs: x})
@@ -251,7 +251,7 @@ def test_entropy_estimates(self):
     _, likelihood = layer(inputs, training=False)
     disc_entropy = tf.reduce_sum(tf.log(likelihood)) / -np.log(2)
     bitstrings = layer.compress(inputs)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       self.assertEqual(1, len(layer.updates))
       sess.run(layer.updates[0])
@@ -272,7 +272,7 @@ def test_noise(self):
     scale = tf.placeholder(tf.float32, [None])
     layer = self.subclass(scale, [1])
     noisy, _ = layer(inputs, training=True)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       values = np.linspace(-50, 50, 100)
       noisy, = sess.run([noisy], {
@@ -288,7 +288,7 @@ def test_quantization(self):
     scale = tf.placeholder(tf.float32, [None])
     layer = self.subclass(scale, [1], mean=None)
     quantized, _ = layer(inputs, training=False)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       values = np.linspace(-50, 50, 100)
       quantized, = sess.run([quantized], {
@@ -305,7 +305,7 @@ def test_quantization_mean(self):
     mean = tf.placeholder(tf.float32, [None])
     layer = self.subclass(scale, [1], mean=mean)
     quantized, _ = layer(inputs, training=False)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       values = np.linspace(-50, 50, 100)
       mean_values = np.random.normal(size=values.shape)
@@ -327,7 +327,7 @@ def test_codec(self):
         scale, [2 ** x for x in range(-10, 10)], mean=None)
     bitstrings = layer.compress(inputs)
     decoded = layer.decompress(bitstrings)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       values = np.linspace(-50, 50, 100)[None]
       decoded, = sess.run([decoded], {
@@ -346,7 +346,7 @@ def test_codec_mean(self):
         scale, [2 ** x for x in range(-10, 10)], mean=mean)
     bitstrings = layer.compress(inputs)
     decoded = layer.decompress(bitstrings)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       values = np.linspace(-50, 50, 100)[None]
       mean_values = np.random.normal(size=values.shape)
@@ -369,7 +369,7 @@ def test_multiple_dimensions(self):
     quantized, _ = layer(inputs, training=False)
     bitstrings = layer.compress(inputs)
     decoded = layer.decompress(bitstrings)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       values = 10 * np.random.normal(size=(2, 5, 3, 7))
       noisy, quantized, decoded = sess.run(
@@ -391,7 +391,7 @@ def test_compress(self):
     layer = self.subclass(scale, scale_table, indexes=indexes)
     bitstrings = layer.compress(inputs)
     decoded = layer.decompress(bitstrings)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       values = 8 * np.random.uniform(size=shape) - 4
       indexes = np.random.randint(
           0, len(scale_table), size=shape, dtype=np.int32)
@@ -415,7 +415,7 @@ def test_decompress(self):
     layer = self.subclass(
         scale, scale_table, indexes=indexes, dtype=tf.float32)
     decoded = layer.decompress(bitstrings)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       decoded, = sess.run([decoded], {
           bitstrings: self.bitstrings,
@@ -437,7 +437,7 @@ def test_quantile_function(self):
     # Test that quantile function inverts cumulative.
     scale = tf.placeholder(tf.float64, [None])
     layer = self.subclass(scale, [1], dtype=tf.float64)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       quantiles = np.array([1e-5, 1e-2, .1, .5, .6, .8])
       locations = layer._standardized_quantile(quantiles)
@@ -452,7 +452,7 @@ def test_distribution(self):
     scale = tf.placeholder(tf.float32, [None, None])
     layer = self.subclass(scale, [1], scale_bound=0, mean=None)
     _, likelihood = layer(inputs, training=False)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       values = np.arange(-5, 1)[:, None]  # must be integers due to quantization
       scales = 2 ** np.linspace(-3, 3, 10)[None, :]
@@ -476,7 +476,7 @@ def test_entropy_estimates(self):
     disc_entropy = tf.reduce_mean(tf.log(likelihood), axis=1)
     disc_entropy /= -np.log(2)
     bitstrings = layer.compress(inputs)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       scales = np.repeat([layer.scale_table], 10000, axis=0).T
       values = self.scipy_class.rvs(scale=scales, size=scales.shape)

tensorflow_compression/python/layers/gdn.py

@@ -158,11 +158,11 @@ def build(self, input_shape):
     # pylint:disable=not-callable
     self.beta = self.beta_parameterizer(
         name="beta", shape=[num_channels], dtype=self.dtype,
-        getter=self.add_variable, initializer=tf.initializers.ones())
+        getter=self.add_weight, initializer=tf.initializers.ones())
 
     self.gamma = self.gamma_parameterizer(
         name="gamma", shape=[num_channels, num_channels], dtype=self.dtype,
-        getter=self.add_variable,
+        getter=self.add_weight,
         initializer=tf.initializers.identity(gain=self._gamma_init))
     # pylint:enable=not-callable
 

tensorflow_compression/python/layers/gdn_test.py

@@ -31,7 +31,7 @@ def _run_gdn(self, x, shape, inverse, rectify, data_format):
     layer = gdn.GDN(
         inverse=inverse, rectify=rectify, data_format=data_format)
     outputs = layer(inputs)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       tf.global_variables_initializer().run()
       y, = sess.run([outputs], {inputs: x})
     return y

tensorflow_compression/python/layers/parameterizers_test.py

@@ -30,7 +30,7 @@ def _test_parameterizer(self, param, init, shape):
     var = param(
         getter=tf.get_variable, name="test", shape=shape, dtype=tf.float32,
         initializer=init, regularizer=None)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       tf.global_variables_initializer().run()
       var, = sess.run([var])
     return var

tensorflow_compression/python/layers/signal_conv.py

@@ -417,10 +417,10 @@ def build(self, input_shape):
 
     kernel_parameterizer = self.kernel_parameterizer
     if kernel_parameterizer is None:
-      getter = self.add_variable
+      getter = self.add_weight
     else:
       getter = functools.partial(
-          kernel_parameterizer, getter=self.add_variable)
+          kernel_parameterizer, getter=self.add_weight)
     self._kernel = getter(
         name="kernel", shape=kernel_shape, dtype=self.dtype,
         initializer=self.kernel_initializer,
@@ -429,10 +429,10 @@ def build(self, input_shape):
     if self.use_bias:
       bias_parameterizer = self.bias_parameterizer
       if bias_parameterizer is None:
-        getter = self.add_variable
+        getter = self.add_weight
       else:
         getter = functools.partial(
-            bias_parameterizer, getter=self.add_variable)
+            bias_parameterizer, getter=self.add_weight)
       self._bias = getter(
           name="bias", shape=(output_channels,), dtype=self.dtype,
           initializer=self.bias_initializer, regularizer=self.bias_regularizer)

tensorflow_compression/python/layers/signal_conv_test.py

@@ -110,7 +110,7 @@ def run_valid(self, batch, input_support, channels, filters, kernel_support,
         activation=activation, use_bias=use_bias,
         kernel_parameterizer=tf_kernel)
     tf_outputs = layer(tf_inputs)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       outputs = sess.run(tf_outputs)
 
@@ -160,7 +160,7 @@ def run_same(self, batch, input_support, channels, filters, kernel_support,
         activation=activation, use_bias=use_bias,
         kernel_parameterizer=tf_kernel)
     tf_outputs = layer(tf_inputs)
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       sess.run(tf.global_variables_initializer())
       outputs = sess.run(tf_outputs)
 

tensorflow_compression/python/ops/math_ops_test.py

@@ -43,7 +43,7 @@ def _test_upper_bound(self, gradient):
       pgrads_expected = [1, 1]
       ngrads_expected = [-1, 0]
 
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       outputs, pgrads, ngrads = sess.run(
           [outputs, pgrads, ngrads], {inputs: inputs_feed})
       self.assertAllEqual(outputs, outputs_expected)
@@ -81,7 +81,7 @@ def _test_lower_bound(self, gradient):
       pgrads_expected = [0, 1]
       ngrads_expected = [-1, -1]
 
-    with self.test_session() as sess:
+    with self.cached_session() as sess:
       outputs, pgrads, ngrads = sess.run(
           [outputs, pgrads, ngrads], {inputs: inputs_feed})
       self.assertAllEqual(outputs, outputs_expected)

tensorflow_compression/python/ops/padding_ops_test.py

@@ -37,7 +37,7 @@ def test_same_padding_corr(self):
             tf.reshape(inputs, (1, 1, -1, 1)),
             tf.reshape(kernel, (1, -1, 1, 1)),
             padding="VALID", data_format="NHWC")
-        with self.test_session() as sess:
+        with self.cached_session() as sess:
           outputs = np.squeeze(sess.run(outputs))
         pos_inp = np.squeeze(np.nonzero(inputs))
         pos_out = np.squeeze(np.nonzero(outputs))
@@ -57,7 +57,7 @@ def test_same_padding_conv(self):
             (1, 1, ishape[0] + kshape[0] - 1, 1),
             strides=(1, 1, 1, 1), padding="VALID", data_format="NHWC")
         outputs = outputs[:, :, (kshape[0] - 1):-(kshape[0] - 1), :]
-        with self.test_session() as sess:
+        with self.cached_session() as sess:
           outputs = np.squeeze(sess.run(outputs))
         pos_inp = np.squeeze(np.nonzero(inputs))
         pos_out = np.squeeze(np.nonzero(outputs))