Changes to pruning to make code compatible with both TF 1.X and TF 2.0 pip

packages by relying on tf.compat.v1. This is to prepare for tf-nightly
switching to 2.0 by default.

Switch some imports to using tf. API, including one that is causing a test to fail
on TF 1.14.

Note that tf.compat.v1 is only available in TF 1.14+, which matches what tf-mot supports.

PiperOrigin-RevId: 265127228

Author: alanchiao

tensorflow_model_optimization/python/core/sparsity/keras/prune_test.py

@@ -18,11 +18,12 @@
 from absl.testing import parameterized
 import numpy as np
 
-from tensorflow.python import keras
-from tensorflow.python.framework import errors_impl
+import tensorflow.compat.v1 as tf
+keras = tf.keras
+errors_impl = tf.errors
+layers = keras.layers
+test = tf.test
 from tensorflow.python.keras import keras_parameterized
-from tensorflow.python.keras import layers
-from tensorflow.python.platform import test
 from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
 from tensorflow_model_optimization.python.core.sparsity.keras import prune
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_callbacks
@@ -349,4 +350,5 @@ def testStripPruningFunctionalModel(self):
 
 
 if __name__ == '__main__':
+  tf.disable_v2_behavior()
   test.main()

tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl_test.py

@@ -19,18 +19,14 @@
 from __future__ import print_function
 # import g3
 import numpy as np
-from tensorflow.python.eager import context
-from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import ops
+
+import tensorflow.compat.v1 as tf
+# TODO(tf-mot): when migrating to 2.0, K.get_session() no longer exists.
+K = tf.keras.backend
+dtypes = tf.dtypes
+test = tf.test
+
 from tensorflow.python.framework import test_util as tf_test_util
-from tensorflow.python.keras import backend as K
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import partitioned_variables
-from tensorflow.python.ops import state_ops
-from tensorflow.python.ops import variable_scope
-from tensorflow.python.ops import variables
-from tensorflow.python.platform import test
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_impl
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils
@@ -66,7 +62,7 @@ def testUpdateSingleMask(self):
     mask_before_pruning = K.get_value(mask)
     self.assertAllEqual(np.count_nonzero(mask_before_pruning), 100)
 
-    if context.executing_eagerly():
+    if tf.executing_eagerly():
       p.conditional_mask_update()
     else:
       K.get_session().run(p.conditional_mask_update())
@@ -121,7 +117,7 @@ def testBlockMaskingAvg(self):
   def testBlockMaskingMax(self):
     block_size = (2, 2)
     block_pooling_type = "MAX"
-    weight = constant_op.constant([[0.1, 0.0, 0.2, 0.0], [0.0, -0.1, 0.0, -0.2],
+    weight = tf.constant([[0.1, 0.0, 0.2, 0.0], [0.0, -0.1, 0.0, -0.2],
                                    [0.3, 0.0, 0.4, 0.0], [0.0, -0.3, 0.0,
                                                           -0.4]])
     expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],
@@ -133,7 +129,7 @@ def testBlockMaskingWithHigherDimensionsRaisesError(self):
     block_size = (2, 2)
     block_pooling_type = "AVG"
     # Weights as in testBlockMasking, but with one extra dimension.
-    weight = constant_op.constant([[[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],
+    weight = tf.constant([[[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],
                                     [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4,
                                                            0.4]]])
     expected_mask = [[[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],
@@ -149,9 +145,9 @@ def testConditionalMaskUpdate(self):
     threshold = K.zeros([])
 
     def linear_sparsity(step):
-      sparsity_val = ops.convert_to_tensor(
+      sparsity_val = tf.convert_to_tensor(
           [0.0, 0.1, 0.1, 0.3, 0.3, 0.5, 0.5, 0.5, 0.5, 0.5])
-      return ops.convert_to_tensor(True), sparsity_val[step]
+      return tf.convert_to_tensor(True), sparsity_val[step]
 
     # Set up pruning
     p = pruning_impl.Pruning(
@@ -163,14 +159,14 @@ def linear_sparsity(step):
 
     non_zero_count = []
     for _ in range(10):
-      if context.executing_eagerly():
+      if tf.executing_eagerly():
         p.conditional_mask_update()
         p.weight_mask_op()
-        state_ops.assign_add(self.global_step, 1)
+        tf.assign_add(self.global_step, 1)
       else:
         K.get_session().run(p.conditional_mask_update())
         K.get_session().run(p.weight_mask_op())
-        K.get_session().run(state_ops.assign_add(self.global_step, 1))
+        K.get_session().run(tf.assign_add(self.global_step, 1))
 
       non_zero_count.append(np.count_nonzero(K.get_value(weight)))
 
@@ -180,4 +176,5 @@ def linear_sparsity(step):
 
 
 if __name__ == "__main__":
+  tf.disable_v2_behavior()
   test.main()

tensorflow_model_optimization/python/core/sparsity/keras/pruning_utils_test.py

@@ -21,26 +21,20 @@
 # import g3
 from absl.testing import parameterized
 
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import nn_ops
-from tensorflow.python.ops import random_ops
-from tensorflow.python.ops import variable_scope
-from tensorflow.python.ops import variables
-from tensorflow.python.platform import test
+import tensorflow.compat.v1 as tf
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils
 
-
 @parameterized.named_parameters(
     ("1x1", [1, 1]), ("4x4", [4, 4]), ("6x6", [6, 6]), ("1x4", [1, 4]),
     ("4x1", [4, 1]), ("1x8", [1, 8]), ("8x1", [8, 1]))
-class PruningUtilsParameterizedTest(test.TestCase, parameterized.TestCase):
+class PruningUtilsParameterizedTest(tf.test.TestCase, parameterized.TestCase):
 
   def _compare_pooling_methods(self, weights, pooling_kwargs):
     with self.cached_session():
-      variables.global_variables_initializer().run()
-      pooled_weights_tf = array_ops.squeeze(
-          nn_ops.pool(
-              array_ops.reshape(
+      tf.global_variables_initializer().run()
+      pooled_weights_tf = tf.squeeze(
+          tf.nn.pool(
+              tf.reshape(
                   weights,
                   [1, weights.get_shape()[0],
                    weights.get_shape()[1], 1]), **pooling_kwargs))
@@ -51,16 +45,16 @@ def _compare_pooling_methods(self, weights, pooling_kwargs):
 
   def _compare_expand_tensor_with_kronecker_product(self, tensor, block_dim):
     with self.cached_session() as session:
-      variables.global_variables_initializer().run()
+      tf.global_variables_initializer().run()
       expanded_tensor = pruning_utils.expand_tensor(tensor, block_dim)
       kronecker_product = pruning_utils.kronecker_product(
-          tensor, array_ops.ones(block_dim))
+          tensor, tf.ones(block_dim))
       expanded_tensor_val, kronecker_product_val = session.run(
           [expanded_tensor, kronecker_product])
       self.assertAllEqual(expanded_tensor_val, kronecker_product_val)
 
   def testFactorizedAvgPool(self, window_shape):
-    weights = variable_scope.get_variable("weights", shape=[1024, 2048])
+    weights = tf.get_variable("weights", shape=[1024, 2048])
     pooling_kwargs = {
         "window_shape": window_shape,
         "pooling_type": "AVG",
@@ -70,7 +64,7 @@ def testFactorizedAvgPool(self, window_shape):
     self._compare_pooling_methods(weights, pooling_kwargs)
 
   def testFactorizedMaxPool(self, window_shape):
-    weights = variable_scope.get_variable("weights", shape=[1024, 2048])
+    weights = tf.get_variable("weights", shape=[1024, 2048])
     pooling_kwargs = {
         "window_shape": window_shape,
         "pooling_type": "MAX",
@@ -80,9 +74,10 @@ def testFactorizedMaxPool(self, window_shape):
     self._compare_pooling_methods(weights, pooling_kwargs)
 
   def testExpandTensor(self, block_dim):
-    weights = random_ops.random_normal(shape=[1024, 512])
+    weights = tf.random.normal(shape=[1024, 512])
     self._compare_expand_tensor_with_kronecker_product(weights, block_dim)
 
 
 if __name__ == "__main__":
-  test.main()
+  tf.disable_v2_behavior()
+  tf.test.main()