Merge pull request #831 from SaoirseARM:toupstream/cluster_default_change

PiperOrigin-RevId: 399130358

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/clustering/keras/BUILD

@@ -25,6 +25,7 @@ py_strict_library(
     srcs_version = "PY3",
     visibility = ["//visibility:public"],
     deps = [
+        ":cluster_config",
         ":cluster_wrapper",
         ":clustering_centroids",
         # tensorflow dep1,

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

@@ -16,11 +16,13 @@
 
 import tensorflow as tf
 
+from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_wrapper
 from tensorflow_model_optimization.python.core.clustering.keras import clustering_centroids
 
 k = tf.keras.backend
 CustomObjectScope = tf.keras.utils.CustomObjectScope
+CentroidInitialization = cluster_config.CentroidInitialization
 Layer = tf.keras.layers.Layer
 InputLayer = tf.keras.layers.InputLayer
 
@@ -44,17 +46,14 @@ def cluster_scope():
     loaded_model = tf.keras.models.load_model(keras_file)
   ```
   """
-  return CustomObjectScope(
-      {
-          'ClusterWeights': cluster_wrapper.ClusterWeights
-      }
-  )
+  return CustomObjectScope({'ClusterWeights': cluster_wrapper.ClusterWeights})
 
 
-def cluster_weights(to_cluster,
-                    number_of_clusters,
-                    cluster_centroids_init,
-                    **kwargs):
+def cluster_weights(
+    to_cluster,
+    number_of_clusters,
+    cluster_centroids_init=CentroidInitialization.KMEANS_PLUS_PLUS,
+    **kwargs):
   """Modifies a keras layer or model to be clustered during training.
 
   This function wraps a keras model or layer with clustering functionality

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

@@ -40,7 +40,6 @@ def setUp(self):
     super(ClusterIntegrationTest, self).setUp()
     self.params = {
         "number_of_clusters": 8,
-        "cluster_centroids_init": CentroidInitialization.LINEAR,
     }
 
     self.x_train = np.array(
@@ -135,6 +134,20 @@ def _get_number_of_unique_weights(self, stripped_model, layer_nr,
     nr_of_unique_weights = len(set(weights_as_list))
     return nr_of_unique_weights
 
+  def testDefaultClusteringInit(self):
+    """Verifies that default initialization method is KMEANS_PLUS_PLUS."""
+    original_model = keras.Sequential([
+        layers.Dense(5, input_shape=(5,)),
+        layers.Dense(5),
+    ])
+    clustered_model = cluster.cluster_weights(original_model, **self.params)
+
+    for i in range(len(clustered_model.layers)):
+      if hasattr(clustered_model.layers[i], "layer"):
+        init_method = clustered_model.layers[i].get_config(
+        )["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."""
@@ -179,7 +192,6 @@ def testSparsityIsPreservedDuringTraining(self):
     original_model.layers[0].set_weights(first_layer_weights)
     clustering_params = {
         "number_of_clusters": 6,
-        "cluster_centroids_init": CentroidInitialization.LINEAR,
         "preserve_sparsity": True
     }
     clustered_model = experimental_cluster.cluster_weights(
@@ -352,7 +364,6 @@ def setUp(self):
 
     self.params_clustering = {
         "number_of_clusters": 16,
-        "cluster_centroids_init": CentroidInitialization.KMEANS_PLUS_PLUS,
     }
 
   def _train(self, model):

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

@@ -54,7 +54,7 @@ class ClusterWeights(Wrapper):
   def __init__(self,
                layer,
                number_of_clusters,
-               cluster_centroids_init,
+               cluster_centroids_init=CentroidInitialization.KMEANS_PLUS_PLUS,
                preserve_sparsity=False,
                cluster_gradient_aggregation=GradientAggregation.SUM,
                **kwargs):

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

@@ -61,7 +61,6 @@ def testCannotBeInitializedWithNonLayerObject(self):
       cluster_wrapper.ClusterWeights(
           {'this': 'is not a Layer instance'},
           number_of_clusters=13,
-          cluster_centroids_init=CentroidInitialization.LINEAR
       )
 
   def testCannotBeInitializedWithNonClusterableLayer(self):
@@ -70,15 +69,13 @@ def testCannotBeInitializedWithNonClusterableLayer(self):
       cluster_wrapper.ClusterWeights(
           NonClusterableLayer(10),
           number_of_clusters=13,
-          cluster_centroids_init=CentroidInitialization.LINEAR
       )
 
   def testCanBeInitializedWithClusterableLayer(self):
     """Verifies that ClusterWeights can be initialized with a built-in clusterable layer."""
     l = cluster_wrapper.ClusterWeights(
         layers.Dense(10),
         number_of_clusters=13,
-        cluster_centroids_init=CentroidInitialization.LINEAR
     )
     self.assertIsInstance(l, cluster_wrapper.ClusterWeights)
 
@@ -87,16 +84,14 @@ def testCannotBeInitializedWithNonIntegerNumberOfClusters(self):
     with self.assertRaises(ValueError):
       cluster_wrapper.ClusterWeights(
           layers.Dense(10),
-          number_of_clusters='13',
-          cluster_centroids_init=CentroidInitialization.LINEAR)
+          number_of_clusters='13')
 
   def testCannotBeInitializedWithFloatNumberOfClusters(self):
     """Verifies that ClusterWeights cannot be initialized with a decimal value provided for the number of clusters."""
     with self.assertRaises(ValueError):
       cluster_wrapper.ClusterWeights(
           layers.Dense(10),
           number_of_clusters=13.4,
-          cluster_centroids_init=CentroidInitialization.LINEAR
       )
 
   @parameterized.parameters(
@@ -111,7 +106,6 @@ def testCannotBeInitializedWithNumberOfClustersLessThanTwo(
       cluster_wrapper.ClusterWeights(
           layers.Dense(10),
           number_of_clusters=number_of_clusters,
-          cluster_centroids_init=CentroidInitialization.LINEAR
       )
 
   @parameterized.parameters((0), (2), (-32))
@@ -122,7 +116,6 @@ def testCannotBeInitializedWithSparsityPreservationAndNumberOfClustersLessThanTh
       cluster_wrapper.ClusterWeights(
           layers.Dense(10),
           number_of_clusters=number_of_clusters,
-          cluster_centroids_init=CentroidInitialization.LINEAR,
           preserve_sparsity=True)
 
   def testCanBeInitializedWithAlreadyClusterableLayer(self):
@@ -131,7 +124,6 @@ def testCanBeInitializedWithAlreadyClusterableLayer(self):
     l = cluster_wrapper.ClusterWeights(
         layer,
         number_of_clusters=13,
-        cluster_centroids_init=CentroidInitialization.LINEAR
     )
     self.assertIsInstance(l, cluster_wrapper.ClusterWeights)
 
@@ -140,7 +132,6 @@ def testIfLayerHasBatchShapeClusterWeightsMustHaveIt(self):
     l = cluster_wrapper.ClusterWeights(
         layers.Dense(10, input_shape=(10,)),
         number_of_clusters=13,
-        cluster_centroids_init=CentroidInitialization.LINEAR
     )
     self.assertTrue(hasattr(l, '_batch_input_shape'))
 
@@ -183,7 +174,6 @@ def testClusterReassociation(self):
     l = cluster_wrapper.ClusterWeights(
         keras.layers.Dense(8, input_shape=input_shape),
         number_of_clusters=2,
-        cluster_centroids_init=CentroidInitialization.LINEAR
     )
     # Build a layer with the given shape
     l.build(input_shape)
@@ -239,7 +229,6 @@ def testSameWeightsAreReturnedBeforeAndAfterSerialisation(self):
     original_layer = cluster_wrapper.ClusterWeights(
         keras.layers.Dense(8, input_shape=input_shape),
         number_of_clusters=2,
-        cluster_centroids_init=CentroidInitialization.LINEAR
     )
     # Build a layer with the given shape
     original_layer.build(input_shape)

tensorflow_model_optimization/python/core/clustering/keras/experimental/BUILD

@@ -24,5 +24,6 @@ py_strict_library(
     visibility = ["//visibility:public"],
     deps = [
         "//tensorflow_model_optimization/python/core/clustering/keras:cluster",
+        "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config",
     ],
 )

tensorflow_model_optimization/python/core/clustering/keras/experimental/cluster.py

@@ -14,11 +14,18 @@
 # ==============================================================================
 """Experimental clustering API functions for Keras models."""
 
+from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 from tensorflow_model_optimization.python.core.clustering.keras.cluster import _cluster_weights
 
+CentroidInitialization = cluster_config.CentroidInitialization
 
-def cluster_weights(to_cluster, number_of_clusters, cluster_centroids_init,
-                    preserve_sparsity, **kwargs):
+
+def cluster_weights(
+    to_cluster,
+    number_of_clusters,
+    cluster_centroids_init=CentroidInitialization.KMEANS_PLUS_PLUS,
+    preserve_sparsity=False,
+    **kwargs):
   """Modify a keras layer or model to be clustered during training (experimental).
 
   This function wraps a keras model or layer with clustering functionality