Replace tensorflow.python.keras with keras. tensorflow.python.keras is an old copy and is deprecated.

PiperOrigin-RevId: 485943740

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -51,8 +51,8 @@ py_strict_library(
     visibility = ["//visibility:public"],
     deps = [
         ":prunable_layer",
+        # keras/engine:base_layer dep1,
         # tensorflow dep1,
-        # python/keras:base_layer tensorflow dep2,
     ],
 )
 
@@ -104,9 +104,9 @@ py_strict_library(
         ":pruning_impl",
         ":pruning_schedule",
         ":pruning_utils",
+        # keras/utils:generic_utils dep1,
         # numpy dep1,
-        # tensorflow dep1,
-        # python/keras/utils:generic_utils tensorflow dep2,
+        # tensorflow:tensorflow_no_contrib dep1,
         "//tensorflow_model_optimization/python/core/keras:compat",
         "//tensorflow_model_optimization/python/core/keras:metrics",
         "//tensorflow_model_optimization/python/core/keras:utils",

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

@@ -21,7 +21,6 @@
 import tensorflow as tf
 
 # TODO(b/139939526): move to public API.
-from tensorflow.python.keras import keras_parameterized
 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 prune_registry
@@ -37,7 +36,6 @@
 ModelCompare = keras_test_utils.ModelCompare
 
 
-@keras_parameterized.run_all_keras_modes
 class PruneIntegrationTest(tf.test.TestCase, parameterized.TestCase,
                            ModelCompare):
 
@@ -691,7 +689,6 @@ def testPruneWithPolynomialDecayPastEndStep_PreservesSparsity(
     self._check_strip_pruning_matches_original(model, 0.6)
 
 
-@keras_parameterized.run_all_keras_modes(always_skip_v1=True)
 class PruneIntegrationCustomTrainingLoopTest(tf.test.TestCase,
                                              parameterized.TestCase):
 

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

@@ -14,12 +14,13 @@
 # ==============================================================================
 """Registry responsible for built-in keras classes."""
 
+from keras.engine import base_layer
 import tensorflow as tf
 
-# TODO(b/139939526): move to public API.
-from tensorflow.python.keras.engine.base_layer import TensorFlowOpLayer
 from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
 
+# TODO(b/139939526): move to public API.
+
 layers = tf.keras.layers
 layers_compat_v1 = tf.compat.v1.keras.layers
 
@@ -100,7 +101,7 @@ class PruneRegistry(object):
       ],
       layers.experimental.SyncBatchNormalization: [],
       layers.experimental.preprocessing.Rescaling.__class__: [],
-      TensorFlowOpLayer: [],
+      base_layer.TensorFlowOpLayer: [],
       layers_compat_v1.BatchNormalization: [],
   }
 

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

@@ -17,12 +17,10 @@
 import json
 import tempfile
 
-from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
 # TODO(b/139939526): move to public API.
-from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.keras import test_utils as keras_test_utils
 from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
 from tensorflow_model_optimization.python.core.sparsity.keras import prune
@@ -57,7 +55,7 @@ class CustomNonPrunableLayer(layers.Dense):
   pass
 
 
-class PruneTest(test.TestCase, parameterized.TestCase):
+class PruneTest(test.TestCase):
 
   INVALID_TO_PRUNE_PARAM_ERROR = ('`prune_low_magnitude` can only prune an '
                                   'object of the following types: '
@@ -202,7 +200,6 @@ def testPruneValidLayersListSuccessful(self):
     for layer, pruned_layer in zip(model_layers, pruned_layers):
       self._validate_pruned_layer(layer, pruned_layer)
 
-  @keras_parameterized.run_all_keras_modes
   def testPruneInferenceWorks_PruningStepCallbackNotRequired(self):
     model = prune.prune_low_magnitude(
         keras.Sequential([

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

@@ -17,12 +17,10 @@
 import os
 import tempfile
 
-from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
 # TODO(b/139939526): move to public API.
-from tensorflow.python.keras import keras_parameterized
 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
@@ -31,7 +29,7 @@
 errors_impl = tf.errors
 
 
-class PruneCallbacksTest(tf.test.TestCase, parameterized.TestCase):
+class PruneCallbacksTest(tf.test.TestCase):
 
   _BATCH_SIZE = 20
 
@@ -55,7 +53,6 @@ def _pruned_model_setup(self, custom_training_loop=False):
       pruned_model.compile(loss=loss, optimizer=optimizer, metrics=['accuracy'])
       return pruned_model, x_train, y_train
 
-  @keras_parameterized.run_all_keras_modes
   def testUpdatePruningStepsAndLogsSummaries(self):
     log_dir = tempfile.mkdtemp()
     pruned_model, x_train, y_train = self._pruned_model_setup()
@@ -77,7 +74,6 @@ def testUpdatePruningStepsAndLogsSummaries(self):
     self._assertLogsExist(log_dir)
 
   # This style of custom training loop isn't available in graph mode.
-  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
   def testUpdatePruningStepsAndLogsSummaries_CustomTrainingLoop(self):
     log_dir = tempfile.mkdtemp()
     pruned_model, loss, optimizer, x_train, y_train = self._pruned_model_setup(
@@ -116,7 +112,6 @@ def testUpdatePruningStepsAndLogsSummaries_CustomTrainingLoop(self):
         3, tf.keras.backend.get_value(pruned_model.layers[1].pruning_step))
     self._assertLogsExist(log_dir)
 
-  @keras_parameterized.run_all_keras_modes
   def testUpdatePruningStepsAndLogsSummaries_RunInference(self):
     pruned_model, _, _, x_train, _ = self._pruned_model_setup(
         custom_training_loop=True)
@@ -128,7 +123,6 @@ def testUpdatePruningStepsAndLogsSummaries_RunInference(self):
     self.assertEqual(
         -1, tf.keras.backend.get_value(pruned_model.layers[1].pruning_step))
 
-  @keras_parameterized.run_all_keras_modes
   def testPruneTrainingRaisesError_PruningStepCallbackMissing(self):
     pruned_model, x_train, y_train = self._pruned_model_setup()
 
@@ -137,7 +131,6 @@ def testPruneTrainingRaisesError_PruningStepCallbackMissing(self):
       pruned_model.fit(x_train, y_train)
 
   # This style of custom training loop isn't available in graph mode.
-  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
   def testPruneTrainingLoopRaisesError_PruningStepCallbackMissing_CustomTrainingLoop(
       self):
     pruned_model, _, _, x_train, _ = self._pruned_model_setup(
@@ -149,7 +142,6 @@ def testPruneTrainingLoopRaisesError_PruningStepCallbackMissing_CustomTrainingLo
       with tf.GradientTape():
         pruned_model(inp, training=True)
 
-  @keras_parameterized.run_all_keras_modes
   def testPruningSummariesRaisesError_LogDirNotNonEmptyString(self):
     with self.assertRaises(ValueError):
       pruning_callbacks.PruningSummaries(log_dir='')

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

@@ -25,11 +25,9 @@
 import tensorflow as tf
 
 # TODO(b/139939526): move to public API.
-from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.keras import compat
 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
 
 K = tf.keras.backend
 dtypes = tf.dtypes
@@ -43,7 +41,6 @@ def assign_add(ref, value):
     return ref.assign_add(value)
 
 
-@keras_parameterized.run_all_keras_modes
 class PruningTest(test.TestCase, parameterized.TestCase):
 
   def setUp(self):

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

@@ -18,7 +18,6 @@
 import tensorflow as tf
 
 # TODO(b/139939526): move to public API.
-from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.keras import compat
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
 
@@ -155,7 +154,6 @@ def testSparsityValueIsValid(self, schedule_type):
 
   # Tests to ensure begin_step, end_step, frequency are used correctly.
 
-  @keras_parameterized.run_all_keras_modes
   @parameterized.named_parameters(
       {
           'testcase_name': 'ConstantSparsity',
@@ -189,7 +187,6 @@ def testPrunesOnlyInBeginEndStepRange(self, schedule_type):
     self.assertFalse(self.evaluate(sparsity(step_201))[0])
     self.assertFalse(self.evaluate(sparsity(step_210))[0])
 
-  @keras_parameterized.run_all_keras_modes
   @parameterized.named_parameters(
       {
           'testcase_name': 'ConstantSparsity',
@@ -216,7 +213,6 @@ def testOnlyPrunesAtValidFrequencySteps(self, schedule_type):
 
 class ConstantSparsityTest(tf.test.TestCase, parameterized.TestCase):
 
-  @keras_parameterized.run_all_keras_modes
   def testPrunesForeverIfEndStepIsNegativeOne(self):
     sparsity = pruning_schedule.ConstantSparsity(0.5, 0, -1, 10)
 
@@ -230,7 +226,6 @@ def testPrunesForeverIfEndStepIsNegativeOne(self):
     self.assertAllClose(0.5, self.evaluate(sparsity(step_10000))[1])
     self.assertAllClose(0.5, self.evaluate(sparsity(step_100000000))[1])
 
-  @keras_parameterized.run_all_keras_modes
   def testPrunesWithConstantSparsity(self):
     sparsity = pruning_schedule.ConstantSparsity(0.5, 100, 200, 10)
 
@@ -263,7 +258,6 @@ def testRaisesErrorIfEndStepIsNegative(self):
     with self.assertRaises(ValueError):
       pruning_schedule.PolynomialDecay(0.4, 0.8, 10, -1)
 
-  @keras_parameterized.run_all_keras_modes
   def testPolynomialDecay_PrunesCorrectly(self):
     sparsity = pruning_schedule.PolynomialDecay(0.2, 0.8, 100, 110, 3, 2)
 

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

@@ -20,12 +20,14 @@
 from __future__ import print_function
 
 import inspect
+
 # import g3
+
+from keras.utils import generic_utils
 import numpy as np
 import tensorflow as tf
 
 # TODO(b/139939526): update to use public API.
-from tensorflow.python.keras.utils import generic_utils
 from tensorflow_model_optimization.python.core.keras import compat as tf_compat
 from tensorflow_model_optimization.python.core.keras import metrics
 from tensorflow_model_optimization.python.core.keras import utils