Migrate majority of pruning API to utilize the tf. public API as opposed to internal versions of those public APIs. Make related updates to support/conform to TF 2.X.

Complete usage of the tf. public API is the only way to ensure that 1.X and 2.X behavior are not mixed and prevents breakages from occurring from TensorFlow file location moves.

PiperOrigin-RevId: 285841197

Author: alanchiao

tensorflow_model_optimization/python/core/keras/BUILD

@@ -15,6 +15,5 @@ py_library(
     deps = [
         # numpy dep1,
         # tensorflow dep1,
-        # python/keras tensorflow dep2,
     ],
 )

tensorflow_model_optimization/python/core/keras/test_utils.py

@@ -13,16 +13,16 @@
 # limitations under the License.
 # ==============================================================================
 """Test utilities for generating, saving, and evaluating models."""
+# TODO(tf-mot): dedup and migrate to testing/ directory.
 
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python import keras
-l = keras.layers
+l = tf.keras.layers
 
 
 def build_simple_dense_model():
-  return keras.Sequential([
+  return tf.keras.Sequential([
       l.Dense(8, activation='relu', input_shape=(10,)),
       l.Dense(5, activation='softmax')
   ])

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -28,10 +28,7 @@ py_library(
     deps = [
         ":pruning_schedule",
         ":pruning_wrapper",
-        # python:util tensorflow dep2,
-        # python/keras tensorflow dep2,
-        # python/keras:engine tensorflow dep2,
-        # python/keras:generic_utils tensorflow dep2,
+        # tensorflow dep1,
     ],
 )
 
@@ -53,7 +50,6 @@ py_library(
     deps = [
         ":prunable_layer",
         # tensorflow dep1,
-        # python/keras:layers_base tensorflow dep2,
     ],
 )
 
@@ -64,10 +60,7 @@ py_library(
     visibility = ["//visibility:public"],
     deps = [
         # six dep1,
-        # python:constant_op tensorflow dep2,
-        # python:dtypes tensorflow dep2,
-        # python:framework_ops tensorflow dep2,
-        # python:math_ops tensorflow dep2,
+        # tensorflow dep1,
     ],
 )
 
@@ -81,6 +74,7 @@ py_test(
         ":pruning_schedule",
         # absl/testing:parameterized dep1,
         # tensorflow dep1,
+        # python/keras tensorflow dep2,
     ],
 )
 
@@ -95,16 +89,8 @@ py_library(
         ":pruning_impl",
         ":pruning_schedule",
         # numpy dep1,
-        # python:check_ops tensorflow dep2,
-        # python:control_flow_ops tensorflow dep2,
-        # python:dtypes tensorflow dep2,
-        # python:framework_ops tensorflow dep2,
-        # python:variables tensorflow dep2,
-        # python/keras:backend tensorflow dep2,
-        # python/keras:base_layer tensorflow dep2,
+        # tensorflow dep1,
         # python/keras:generic_utils tensorflow dep2,
-        # python/keras:initializers tensorflow dep2,
-        # python/keras:layers_base tensorflow dep2,
         # python/keras:tf_utils tensorflow dep2,
     ],
 )
@@ -118,9 +104,6 @@ py_library(
         ":pruning_wrapper",
         # numpy dep1,
         # tensorflow dep1,
-        # python:math_ops tensorflow dep2,
-        # python/keras:backend tensorflow dep2,
-        # python/keras:callbacks tensorflow dep2,
     ],
 )
 
@@ -149,15 +132,8 @@ py_library(
     deps = [
         ":pruning_utils",
         # tensorflow dep1,
-        # python:array_ops tensorflow dep2,
-        # python:control_flow_ops tensorflow dep2,
-        # python:dtypes tensorflow dep2,
-        # python:framework_ops tensorflow dep2,
-        # python:math_ops tensorflow dep2,
-        # python:nn_ops tensorflow dep2,
         # python:state_ops tensorflow dep2,
         # python:summary tensorflow dep2,
-        # python:variables tensorflow dep2,
     ],
 )
 
@@ -168,15 +144,7 @@ py_library(
     visibility = ["//visibility:public"],
     deps = [
         # numpy dep1,
-        # python:array_ops tensorflow dep2,
-        # python:constant_op tensorflow dep2,
-        # python:control_flow_ops tensorflow dep2,
-        # python:dtypes tensorflow dep2,
-        # python:framework_ops tensorflow dep2,
-        # python:init_ops tensorflow dep2,
-        # python:nn_ops tensorflow dep2,
-        # python:state_ops tensorflow dep2,
-        # python:variable_scope tensorflow dep2,
+        # tensorflow dep1,
     ],
 )
 
@@ -259,7 +227,6 @@ py_test(
         # absl/testing:parameterized dep1,
         # numpy dep1,
         # tensorflow dep1,
-        # python/keras tensorflow dep2,
         "//tensorflow_model_optimization/python/core/keras:test_utils",
     ],
 )
@@ -300,6 +267,8 @@ py_test(
         ":pruning_impl",
         ":pruning_schedule",
         # numpy dep1,
+        # python/keras tensorflow dep2,
+        "//tensorflow_model_optimization/python/core/keras:test_utils",
     ],
 )
 

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

@@ -15,12 +15,14 @@
 # pylint: disable=protected-access,missing-docstring,unused-argument
 """Entry point for pruning models during training."""
 
-from tensorflow.python import keras
-from tensorflow.python.keras.engine.input_layer import InputLayer
-from tensorflow.python.keras.utils.generic_utils import custom_object_scope
+import tensorflow as tf
+
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule as pruning_sched
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper
 
+keras = tf.keras
+custom_object_scope = tf.keras.utils.custom_object_scope
+
 
 def prune_scope():
   """Provides a scope in which Pruned layers and models can be deserialized.
@@ -125,7 +127,7 @@ def _prune_list(layers, **params):
       # No need to wrap the input layer either.
       if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
         wrapped_layers.append(layer)
-      elif isinstance(layer, InputLayer):
+      elif isinstance(layer, keras.layers.InputLayer):
         # TODO(yunluli): Replace with a clone function in keras.
         wrapped_layers.append(layer.__class__.from_config(layer.get_config()))
       else:

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

@@ -17,31 +17,28 @@
 import tempfile
 from absl.testing import parameterized
 import numpy as np
+import tensorflow as tf
 
-from tensorflow.python import keras
-from tensorflow.python.distribute import collective_all_reduce_strategy
-from tensorflow.python.distribute import mirrored_strategy
-from tensorflow.python.distribute import one_device_strategy
-from tensorflow.python.keras.utils import np_utils
-from tensorflow.python.platform import test
 from tensorflow_model_optimization.python.core.keras import test_utils as keras_test_utils
 from tensorflow_model_optimization.python.core.sparsity.keras import prune
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_callbacks
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
 from tensorflow_model_optimization.python.core.sparsity.keras import test_utils
 
+keras = tf.keras
+
 
 def _distribution_strategies():
   return [
-      collective_all_reduce_strategy.CollectiveAllReduceStrategy(),
-      mirrored_strategy.MirroredStrategy(),
+      tf.distribute.experimental.MultiWorkerMirroredStrategy(),
+      tf.distribute.MirroredStrategy(),
       # TODO(pulkitb): Add parameter_server
-      # parameter_server_strategy.ParameterServerStrategy(),
-      one_device_strategy.OneDeviceStrategy('/cpu:0'),
+      # tf.distribute.experimental.ParameterServerStrategy,
+      tf.distribute.OneDeviceStrategy('/cpu:0'),
   ]
 
 
-class PruneDistributedTest(test.TestCase, parameterized.TestCase):
+class PruneDistributedTest(tf.test.TestCase, parameterized.TestCase):
 
   def setUp(self):
     super(PruneDistributedTest, self).setUp()
@@ -67,8 +64,7 @@ def testPrunesSimpleDenseModel(self, distribution):
     # Simple unpruned model. No sparsity.
     model.fit(
         np.random.rand(20, 10),
-        np_utils.to_categorical(
-            np.random.randint(5, size=(20, 1)), 5),
+        keras.utils.to_categorical(np.random.randint(5, size=(20, 1)), 5),
         epochs=2,
         callbacks=[pruning_callbacks.UpdatePruningStep()],
         batch_size=20)
@@ -85,4 +81,4 @@ def testPrunesSimpleDenseModel(self, distribution):
 
 
 if __name__ == '__main__':
-  test.main()
+  tf.test.main()