Updated Clustering initialization default to KMEANS_PLUS_PLUS

* Added integration test which ensures the default is KMEANS_PLUS_PLUS

Author: SaoirseARM

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

@@ -18,9 +18,11 @@
 
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_wrapper
 from tensorflow_model_optimization.python.core.clustering.keras import clustering_centroids
+from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 
 k = tf.keras.backend
 CustomObjectScope = tf.keras.utils.CustomObjectScope
+CentroidInitialization = cluster_config.CentroidInitialization
 Layer = tf.keras.layers.Layer
 InputLayer = tf.keras.layers.InputLayer
 
@@ -53,7 +55,7 @@ def cluster_scope():
 
 def cluster_weights(to_cluster,
                     number_of_clusters,
-                    cluster_centroids_init,
+                    cluster_centroids_init=CentroidInitialization.KMEANS_PLUS_PLUS,
                     **kwargs):
   """Modifies a keras layer or model to be clustered during training.
 

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,19 @@ 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 +191,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 +363,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/cluster.py

@@ -15,10 +15,12 @@
 """Experimental clustering API functions for Keras models."""
 
 from tensorflow_model_optimization.python.core.clustering.keras.cluster import _cluster_weights
+from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 
+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