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

PiperOrigin-RevId: 485928216

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/clustering/keras/cluster_integration_test.py

@@ -21,7 +21,6 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.clustering.keras import cluster
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 from tensorflow_model_optimization.python.core.clustering.keras.experimental import cluster as experimental_cluster
@@ -153,7 +152,6 @@ def testDefaultClusteringInit(self):
         )["cluster_centroids_init"]
         self.assertEqual(init_method, CentroidInitialization.KMEANS_PLUS_PLUS)
 
-  @keras_parameterized.run_all_keras_modes
   def testValuesRemainClusteredAfterTraining(self):
     """Verifies that training a clustered model does not destroy the clusters."""
     original_model = keras.Sequential([
@@ -175,7 +173,6 @@ def testValuesRemainClusteredAfterTraining(self):
     unique_weights = set(weights_as_list)
     self.assertLessEqual(len(unique_weights), self.params["number_of_clusters"])
 
-  @keras_parameterized.run_all_keras_modes
   def testSparsityIsPreservedDuringTraining(self):
     """Set a specific random seed.
 
@@ -230,7 +227,6 @@ def testSparsityIsPreservedDuringTraining(self):
         nr_of_unique_weights_after,
         clustering_params["number_of_clusters"])
 
-  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
   def testEndToEndSequential(self):
     """Test End to End clustering - sequential model."""
     original_model = keras.Sequential([
@@ -247,7 +243,6 @@ def clusters_check(stripped_model):
 
     self.end_to_end_testing(original_model, clusters_check)
 
-  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
   def testEndToEndConv1DAndConv1DTranspose(self):
     """Test End to End clustering - model with Conv1D and Conv1DTranspose."""
     inp = layers.Input(batch_shape=(1, 16))
@@ -372,7 +367,6 @@ def clusters_check(stripped_model):
 
     self.end_to_end_testing(original_model, clusters_check)
 
-  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
   def testEndToEndFunctional(self):
     """Test End to End clustering - functional model."""
     inputs = keras.layers.Input(shape=(5,))
@@ -389,7 +383,6 @@ def clusters_check(stripped_model):
 
     self.end_to_end_testing(original_model, clusters_check)
 
-  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
   def testEndToEndDeepLayer(self):
     """Test End to End clustering for the model with deep layer."""
     internal_model = tf.keras.Sequential(
@@ -416,7 +409,6 @@ def clusters_check(stripped_model):
 
     self.end_to_end_testing(original_model, clusters_check)
 
-  @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
   def testEndToEndDeepLayer2(self):
     """Test End to End clustering for the model with 2 deep layers."""
     internal_model = tf.keras.Sequential(
@@ -454,7 +446,6 @@ def clusters_check(stripped_model):
 
     self.end_to_end_testing(original_model, clusters_check)
 
-  @keras_parameterized.run_all_keras_modes
   def testWeightsAreLearningDuringClustering(self):
     """Verifies that weights are updated during training a clustered model.
 
@@ -541,7 +532,6 @@ def _assertNbUniqueWeights(self, weight, expected_unique_weights):
     nr_unique_weights = len(np.unique(weight.numpy().flatten()))
     assert nr_unique_weights == expected_unique_weights
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterSimpleRNN(self):
     model = keras.models.Sequential()
     model.add(keras.layers.Embedding(self.max_features, 16,
@@ -564,7 +554,6 @@ def testClusterSimpleRNN(self):
 
     self._train(stripped_model)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterLSTM(self):
     model = keras.models.Sequential()
     model.add(keras.layers.Embedding(self.max_features, 16,
@@ -587,7 +576,6 @@ def testClusterLSTM(self):
 
     self._train(stripped_model)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterGRU(self):
     model = keras.models.Sequential()
     model.add(keras.layers.Embedding(self.max_features, 16,
@@ -610,7 +598,6 @@ def testClusterGRU(self):
 
     self._train(stripped_model)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterBidirectional(self):
     model = keras.models.Sequential()
     model.add(
@@ -634,7 +621,6 @@ def testClusterBidirectional(self):
         expected_unique_weights=self.params_clustering["number_of_clusters"],
     )
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterStackedRNNCells(self):
     model = keras.models.Sequential()
     model.add(
@@ -685,7 +671,6 @@ def _get_model(self):
     model = tf.keras.Model(inputs=inp, outputs=out)
     return model
 
-  @keras_parameterized.run_all_keras_modes
   def testMHA(self):
     model = self._get_model()
 
@@ -736,7 +721,6 @@ def _get_model(self):
     model = tf.keras.Model(inputs=inp, outputs=out)
     return model
 
-  @keras_parameterized.run_all_keras_modes
   def testPerChannel(self):
     model = self._get_model()
 

tensorflow_model_optimization/python/core/clustering/keras/cluster_test.py

@@ -21,7 +21,6 @@
 from absl.testing import parameterized
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.clustering.keras import cluster
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_wrapper
@@ -162,15 +161,13 @@ def _count_clustered_layers(self, model):
         count += 1
     return count
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterKerasClusterableLayer(self):
     """Verifies that a built-in keras layer marked as clusterable is being clustered correctly."""
     wrapped_layer = self._build_clustered_layer_model(
         self.keras_clusterable_layer)
 
     self._validate_clustered_layer(self.keras_clusterable_layer, wrapped_layer)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterKerasClusterableLayerWithSparsityPreservation(self):
     """Verifies that a built-in keras layer marked as clusterable is being clustered correctly when sparsity preservation is enabled."""
     preserve_sparsity_params = {'preserve_sparsity': True}
@@ -180,7 +177,6 @@ def testClusterKerasClusterableLayerWithSparsityPreservation(self):
 
     self._validate_clustered_layer(self.keras_clusterable_layer, wrapped_layer)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterKerasNonClusterableLayer(self):
     """Verifies that a built-in keras layer not marked as clusterable is not being clustered."""
     wrapped_layer = self._build_clustered_layer_model(
@@ -190,7 +186,6 @@ def testClusterKerasNonClusterableLayer(self):
                                    wrapped_layer)
     self.assertEqual([], wrapped_layer.layer.get_clusterable_weights())
 
-  @keras_parameterized.run_all_keras_modes
   def testDepthwiseConv2DLayerNonClusterable(self):
     """Verifies that we don't cluster a DepthwiseConv2D layer, because clustering of this type of layer gives big unrecoverable accuracy loss."""
     wrapped_layer = self._build_clustered_layer_model(
@@ -200,7 +195,6 @@ def testDepthwiseConv2DLayerNonClusterable(self):
                                    wrapped_layer)
     self.assertEqual([], wrapped_layer.layer.get_clusterable_weights())
 
-  @keras_parameterized.run_all_keras_modes
   def testDenseLayer(self):
     """Verifies that we can cluster a Dense layer."""
     input_shape = (28, 1)
@@ -214,7 +208,6 @@ def testDenseLayer(self):
     self.assertEqual([1, 10],
                      wrapped_layer.layer.get_clusterable_weights()[0][1].shape)
 
-  @keras_parameterized.run_all_keras_modes
   def testConv1DLayer(self):
     """Verifies that we can cluster a Conv1D layer."""
     input_shape = (28, 1)
@@ -227,7 +220,6 @@ def testConv1DLayer(self):
     self.assertEqual([5, 1, 3],
                      wrapped_layer.layer.get_clusterable_weights()[0][1].shape)
 
-  @keras_parameterized.run_all_keras_modes
   def testConv1DTransposeLayer(self):
     """Verifies that we can cluster a Conv1DTranspose layer."""
     input_shape = (28, 1)
@@ -240,7 +232,6 @@ def testConv1DTransposeLayer(self):
     self.assertEqual([5, 3, 1],
                      wrapped_layer.layer.get_clusterable_weights()[0][1].shape)
 
-  @keras_parameterized.run_all_keras_modes
   def testConv2DLayer(self):
     """Verifies that we can cluster a Conv2D layer."""
     input_shape = (28, 28, 1)
@@ -253,7 +244,6 @@ def testConv2DLayer(self):
     self.assertEqual([4, 5, 1, 3],
                      wrapped_layer.layer.get_clusterable_weights()[0][1].shape)
 
-  @keras_parameterized.run_all_keras_modes
   def testConv2DTransposeLayer(self):
     """Verifies that we can cluster a Conv2DTranspose layer."""
     input_shape = (28, 28, 1)
@@ -266,7 +256,6 @@ def testConv2DTransposeLayer(self):
     self.assertEqual([4, 5, 3, 1],
                      wrapped_layer.layer.get_clusterable_weights()[0][1].shape)
 
-  @keras_parameterized.run_all_keras_modes
   def testConv3DLayer(self):
     """Verifies that we can cluster a Conv3D layer."""
     input_shape = (28, 28, 28, 1)
@@ -287,7 +276,6 @@ def testClusterKerasUnsupportedLayer(self):
     with self.assertRaises(ValueError):
       cluster.cluster_weights(keras_unsupported_layer, **self.params)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterCustomClusterableLayer(self):
     """Verifies that a custom clusterable layer is being clustered correctly."""
     wrapped_layer = self._build_clustered_layer_model(
@@ -297,7 +285,6 @@ def testClusterCustomClusterableLayer(self):
     self.assertEqual([('kernel', wrapped_layer.layer.kernel)],
                      wrapped_layer.layer.get_clusterable_weights())
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterCustomClusterableLayerWithSparsityPreservation(self):
     """Verifies that a custom clusterable layer is being clustered correctly when sparsity preservation is enabled."""
     preserve_sparsity_params = {'preserve_sparsity': True}
@@ -424,7 +411,6 @@ def testStripClusteringSequentialModelWithBiasConstraint(self):
       keras_file = os.path.join(tmp_dir_name, 'cluster_test')
       stripped_model.save(keras_file, save_traces=True)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterSequentialModelSelectively(self):
     clustered_model = keras.Sequential()
     clustered_model.add(
@@ -437,7 +423,6 @@ def testClusterSequentialModelSelectively(self):
     self.assertNotIsInstance(clustered_model.layers[1],
                              cluster_wrapper.ClusterWeights)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterSequentialModelSelectivelyWithSparsityPreservation(self):
     """Verifies that layers within a sequential model can be clustered selectively when sparsity preservation is enabled."""
     preserve_sparsity_params = {'preserve_sparsity': True}
@@ -454,7 +439,6 @@ def testClusterSequentialModelSelectivelyWithSparsityPreservation(self):
     self.assertNotIsInstance(clustered_model.layers[1],
                              cluster_wrapper.ClusterWeights)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterFunctionalModelSelectively(self):
     """Verifies that layers within a functional model can be clustered selectively."""
     i1 = keras.Input(shape=(10,))
@@ -469,7 +453,6 @@ def testClusterFunctionalModelSelectively(self):
     self.assertNotIsInstance(clustered_model.layers[3],
                              cluster_wrapper.ClusterWeights)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterFunctionalModelSelectivelyWithSparsityPreservation(self):
     """Verifies that layers within a functional model can be clustered selectively when sparsity preservation is enabled."""
     preserve_sparsity_params = {'preserve_sparsity': True}
@@ -486,7 +469,6 @@ def testClusterFunctionalModelSelectivelyWithSparsityPreservation(self):
     self.assertNotIsInstance(clustered_model.layers[3],
                              cluster_wrapper.ClusterWeights)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterModelValidLayersSuccessful(self):
     """Verifies that clustering a sequential model results in all clusterable layers within the model being clustered."""
     model = keras.Sequential([
@@ -500,7 +482,6 @@ def testClusterModelValidLayersSuccessful(self):
     for layer, clustered_layer in zip(model.layers, clustered_model.layers):
       self._validate_clustered_layer(layer, clustered_layer)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterModelValidLayersSuccessfulWithSparsityPreservation(self):
     """Verifies that clustering a sequential model results in all clusterable layers within the model being clustered when sparsity preservation is enabled."""
     preserve_sparsity_params = {'preserve_sparsity': True}
@@ -540,7 +521,6 @@ def testClusterModelCustomNonClusterableLayerRaisesError(self):
               self.custom_clusterable_layer, custom_non_clusterable_layer
           ]), **self.params)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterModelDoesNotWrapAlreadyWrappedLayer(self):
     """Verifies that clustering a model that contains an already clustered layer does not result in wrapping the clustered layer into another cluster_wrapper."""
     model = keras.Sequential([
@@ -579,7 +559,6 @@ def testClusterSequentialModelNoInput(self):
     clustered_model = cluster.cluster_weights(model, **self.params)
     self.assertEqual(self._count_clustered_layers(clustered_model), 2)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterSequentialModelWithInput(self):
     """Verifies that a sequential model with an input layer is being clustered correctly."""
     # With InputLayer
@@ -607,7 +586,6 @@ def testClusterSequentialModelPreservesBuiltStateNoInput(self):
           json.loads(clustered_model.to_json()))
       self.assertEqual(loaded_model.built, False)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterSequentialModelPreservesBuiltStateWithInput(self):
     """Verifies that clustering a sequential model with an input layer preserves the built state of the model."""
     # With InputLayer
@@ -625,7 +603,6 @@ def testClusterSequentialModelPreservesBuiltStateWithInput(self):
           json.loads(clustered_model.to_json()))
     self.assertEqual(loaded_model.built, True)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterFunctionalModelPreservesBuiltState(self):
     """Verifies that clustering a functional model preserves the built state of the model."""
     i1 = keras.Input(shape=(10,))
@@ -644,7 +621,6 @@ def testClusterFunctionalModelPreservesBuiltState(self):
           json.loads(clustered_model.to_json()))
     self.assertEqual(loaded_model.built, True)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterFunctionalModel(self):
     """Verifies that a functional model is being clustered correctly."""
     i1 = keras.Input(shape=(10,))
@@ -656,7 +632,6 @@ def testClusterFunctionalModel(self):
     clustered_model = cluster.cluster_weights(model, **self.params)
     self.assertEqual(self._count_clustered_layers(clustered_model), 3)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterFunctionalModelWithLayerReused(self):
     """Verifies that a layer reused within a functional model multiple times is only being clustered once."""
     # The model reuses the Dense() layer. Make sure it's only clustered once.
@@ -668,22 +643,19 @@ def testClusterFunctionalModelWithLayerReused(self):
     clustered_model = cluster.cluster_weights(model, **self.params)
     self.assertEqual(self._count_clustered_layers(clustered_model), 1)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterSubclassModel(self):
     """Verifies that attempting to cluster an instance of a subclass of keras.Model raises an exception."""
     model = TestModel()
     with self.assertRaises(ValueError):
       _ = cluster.cluster_weights(model, **self.params)
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterSubclassModelAsSubmodel(self):
     """Verifies that attempting to cluster a model with submodel that is a subclass throws an exception."""
     model_subclass = TestModel()
     model = keras.Sequential([layers.Dense(10), model_subclass])
     with self.assertRaisesRegex(ValueError, 'Subclassed models.*'):
       _ = cluster.cluster_weights(model, **self.params)
 
-  @keras_parameterized.run_all_keras_modes
   def testStripClusteringSequentialModel(self):
     """Verifies that stripping the clustering wrappers from a sequential model produces the expected config."""
     model = keras.Sequential([
@@ -697,7 +669,6 @@ def testStripClusteringSequentialModel(self):
     self.assertEqual(self._count_clustered_layers(stripped_model), 0)
     self.assertEqual(model.get_config(), stripped_model.get_config())
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterStrippingFunctionalModel(self):
     """Verifies that stripping the clustering wrappers from a functional model produces the expected config."""
     i1 = keras.Input(shape=(10,))
@@ -713,7 +684,6 @@ def testClusterStrippingFunctionalModel(self):
     self.assertEqual(self._count_clustered_layers(stripped_model), 0)
     self.assertEqual(model.get_config(), stripped_model.get_config())
 
-  @keras_parameterized.run_all_keras_modes
   def testClusterWeightsStrippedWeights(self):
     """Verifies that stripping the clustering wrappers from a functional model preserves the clustered weights."""
     i1 = keras.Input(shape=(10,))
@@ -728,7 +698,6 @@ def testClusterWeightsStrippedWeights(self):
     self.assertEqual(self._count_clustered_layers(stripped_model), 0)
     self.assertLen(stripped_model.get_weights(), cluster_weight_length)
 
-  @keras_parameterized.run_all_keras_modes
   def testStrippedKernel(self):
     """Verifies that stripping the clustering wrappers from a functional model restores the layers kernel and the layers weight array to the new clustered weight value."""
     i1 = keras.Input(shape=(1, 1, 1))
@@ -746,7 +715,6 @@ def testStrippedKernel(self):
     self.assertIsNot(stripped_conv2d_layer.kernel, clustered_kernel)
     self.assertIn(stripped_conv2d_layer.kernel, stripped_conv2d_layer.weights)
 
-  @keras_parameterized.run_all_keras_modes
   def testStripSelectivelyClusteredFunctionalModel(self):
     """Verifies that invoking strip_clustering() on a selectively clustered functional model strips the clustering wrappers from the clustered layers."""
     i1 = keras.Input(shape=(10,))
@@ -761,7 +729,6 @@ def testStripSelectivelyClusteredFunctionalModel(self):
     self.assertEqual(self._count_clustered_layers(stripped_model), 0)
     self.assertIsInstance(stripped_model.layers[2], layers.Dense)
 
-  @keras_parameterized.run_all_keras_modes
   def testStripSelectivelyClusteredSequentialModel(self):
     """Verifies that invoking strip_clustering() on a selectively clustered sequential model strips the clustering wrappers from the clustered layers."""
     clustered_model = keras.Sequential([
@@ -775,7 +742,6 @@ def testStripSelectivelyClusteredSequentialModel(self):
     self.assertEqual(self._count_clustered_layers(stripped_model), 0)
     self.assertIsInstance(stripped_model.layers[0], layers.Dense)
 
-  @keras_parameterized.run_all_keras_modes
   def testStripClusteringAndSetOriginalWeightsBack(self):
     """Verifies that we can set_weights onto the stripped model."""
     model = keras.Sequential([