Add Anti Zero-Drift functionality for Sparsity-Aware clustering

 * Implemented the zero-centroid initialization for all clustering methods
 * Implemented the sparsity masks for forward and backward propagation
 * Added preserve_sparsity class member to ClusterWeights to make
   sparsity preservation optional for all clustering methods
 * Refactored AbstractCentroidsInitialisation to include zero-centroid
   initialization for all init types
 * Added unit tests around the new changes

Author: MatteoArm

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

@@ -56,6 +56,7 @@ def cluster_scope():
 def cluster_weights(to_cluster,
                     number_of_clusters,
                     cluster_centroids_init,
+                    preserve_sparsity=False,
                     **kwargs):
   """Modify a keras layer or model to be clustered during training.
 
@@ -108,8 +109,19 @@ def cluster_weights(to_cluster,
       number_of_clusters: the number of cluster centroids to form when
         clustering a layer/model. For example, if number_of_clusters=8 then only
         8 unique values will be used in each weight array.
-      cluster_centroids_init: `tfmot.clustering.keras.CentroidInitialization`
-        instance that determines how the cluster centroids will be initialized.
+      cluster_centroids_init: enum value that determines how the cluster
+        centroids will be initialized.
+        Can have following values:
+          1. RANDOM : centroids are sampled using the uniform distribution
+          between the minimum and maximum weight values in a given layer
+          2. DENSITY_BASED : density-based sampling. First, cumulative
+          distribution function is built for weights, then y-axis is evenly
+          spaced into number_of_clusters regions. After this the corresponding x
+          values are obtained and used to initialize clusters centroids.
+          3. LINEAR : cluster centroids are evenly spaced between the minimum
+          and maximum values of a given weight
+      preserve_sparsity: optional boolean value that determines whether or not
+        sparsity preservation will be enforced during training
       **kwargs: Additional keyword arguments to be passed to the keras layer.
         Ignored when to_cluster is not a keras layer.
 
@@ -146,6 +158,7 @@ def _add_clustering_wrapper(layer):
     return cluster_wrapper.ClusterWeights(layer,
                                           number_of_clusters,
                                           cluster_centroids_init,
+                                          preserve_sparsity,
                                           **kwargs)
 
   def _wrap_list(layers):

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

@@ -60,10 +60,26 @@ def setUp(self):
         dtype="float32",
     )
 
+    self.x_train2 = np.array(
+        [[0.0, 1.0, 2.0, 3.0, 4.0], [2.0, 0.0, 2.0, 3.0, 4.0], [0.0, 3.0, 2.0, 3.0, 4.0],
+         [4.0, 1.0, 2.0, 3.0, 4.0], [5.0, 1.0, 2.0, 3.0, 4.0]],
+        dtype="float32",
+    )
+
+    self.y_train2 = np.array(
+        [[0.0, 1.0, 2.0, 3.0, 4.0], [1.0, 0.0, 2.0, 3.0, 4.0], [1.0, 0.0, 2.0, 3.0, 4.0],
+         [0.0, 1.0, 2.0, 3.0, 4.0], [0.0, 1.0, 2.0, 3.0, 4.0]],
+        dtype="float32",
+    )
+
   def dataset_generator(self):
     for x, y in zip(self.x_train, self.y_train):
       yield np.array([x]), np.array([y])
 
+  def dataset_generator2(self):
+    for x, y in zip(self.x_train2, self.y_train2):
+      yield np.array([x]), np.array([y])
+
   def end_to_end_testing(self, original_model, clusters_check=None):
     """Test End to End clustering."""
 
@@ -128,6 +144,47 @@ 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):
+    """Verifies that training a clustered model does not destroy the sparsity of the weights."""
+    original_model = keras.Sequential([
+        layers.Dense(5, input_shape=(5,)),
+        layers.Dense(5),
+    ])
+
+    """Using a mininum number of centroids to make it more likely that some weights will be zero."""
+    clustering_params = {
+        "number_of_clusters": 3,
+        "cluster_centroids_init": CentroidInitialization.LINEAR,
+        "preserve_sparsity": True
+    }
+
+    clustered_model = cluster.cluster_weights(original_model, **clustering_params)
+
+    stripped_model_before_tuning = cluster.strip_clustering(clustered_model)
+    weights_before_tuning = stripped_model_before_tuning.get_weights()[0]
+    non_zero_weight_indices_before_tuning = np.nonzero(weights_before_tuning)
+
+    clustered_model.compile(
+        loss=keras.losses.categorical_crossentropy,
+        optimizer="adam",
+        metrics=["accuracy"],
+    )
+    clustered_model.fit(x=self.dataset_generator2(), steps_per_epoch=1)
+
+    stripped_model_after_tuning = cluster.strip_clustering(clustered_model)
+    weights_after_tuning = stripped_model_after_tuning.get_weights()[0]
+    non_zero_weight_indices_after_tuning = np.nonzero(weights_after_tuning)
+    weights_as_list_after_tuning = weights_after_tuning.reshape(-1,).tolist()
+    unique_weights_after_tuning = set(weights_as_list_after_tuning)
+
+    """Check that the null weights stayed the same before and after tuning."""
+    self.assertTrue(np.array_equal(non_zero_weight_indices_before_tuning,
+                                   non_zero_weight_indices_after_tuning))
+
+    """Check that the number of unique weights matches the number of clusters."""
+    self.assertLessEqual(len(unique_weights_after_tuning), self.params["number_of_clusters"])
+
   @keras_parameterized.run_all_keras_modes(always_skip_v1=True)
   def testEndToEndSequential(self):
     """Test End to End clustering - sequential model."""

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

@@ -112,6 +112,18 @@ def testClusterKerasClusterableLayer(self):
 
     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 }
+    params = { **self.params, **preserve_sparsity_params }
+    wrapped_layer = cluster.cluster_weights(self.keras_clusterable_layer, **params)
+
+    self._validate_clustered_layer(self.keras_clusterable_layer, wrapped_layer)
+
   @keras_parameterized.run_all_keras_modes
   def testClusterKerasNonClusterableLayer(self):
     """
@@ -164,6 +176,22 @@ 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 }
+    params = { **self.params, **preserve_sparsity_params }
+    wrapped_layer = cluster.cluster_weights(self.custom_clusterable_layer, **params)
+    self.model.add(wrapped_layer)
+    self.model.build(input_shape=(10, 1))
+
+    self._validate_clustered_layer(self.custom_clusterable_layer, wrapped_layer)
+    self.assertEqual([('kernel', wrapped_layer.layer.kernel)],
+                     wrapped_layer.layer.get_clusterable_weights())
+
   def testClusterCustomNonClusterableLayer(self):
     """
     Verifies that attempting to cluster a custom non-clusterable layer raises
@@ -193,6 +221,22 @@ def testClusterSequentialModelSelectively(self):
     self.assertIsInstance(clustered_model.layers[0], cluster_wrapper.ClusterWeights)
     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 }
+    params = { **self.params, **preserve_sparsity_params }
+    clustered_model = keras.Sequential()
+    clustered_model.add(cluster.cluster_weights(self.keras_clusterable_layer, **params))
+    clustered_model.add(self.keras_clusterable_layer)
+    clustered_model.build(input_shape=(1, 10))
+
+    self.assertIsInstance(clustered_model.layers[0], cluster_wrapper.ClusterWeights)
+    self.assertNotIsInstance(clustered_model.layers[1], cluster_wrapper.ClusterWeights)
+
   @keras_parameterized.run_all_keras_modes
   def testClusterFunctionalModelSelectively(self):
     """
@@ -209,6 +253,24 @@ def testClusterFunctionalModelSelectively(self):
     self.assertIsInstance(clustered_model.layers[2], cluster_wrapper.ClusterWeights)
     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 }
+    params = { **self.params, **preserve_sparsity_params }
+    i1 = keras.Input(shape=(10,))
+    i2 = keras.Input(shape=(10,))
+    x1 = cluster.cluster_weights(layers.Dense(10), **params)(i1)
+    x2 = layers.Dense(10)(i2)
+    outputs = layers.Add()([x1, x2])
+    clustered_model = keras.Model(inputs=[i1, i2], outputs=outputs)
+
+    self.assertIsInstance(clustered_model.layers[2], cluster_wrapper.ClusterWeights)
+    self.assertNotIsInstance(clustered_model.layers[3], cluster_wrapper.ClusterWeights)
+
   @keras_parameterized.run_all_keras_modes
   def testClusterModelValidLayersSuccessful(self):
     """
@@ -227,6 +289,26 @@ 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 }
+    params = { **self.params, **preserve_sparsity_params }
+    model = keras.Sequential([
+        self.keras_clusterable_layer,
+        self.keras_non_clusterable_layer,
+        self.custom_clusterable_layer
+    ])
+    clustered_model = cluster.cluster_weights(model, **params)
+    clustered_model.build(input_shape=(1, 28, 28, 1))
+
+    self.assertEqual(len(model.layers), len(clustered_model.layers))
+    for layer, clustered_layer in zip(model.layers, clustered_model.layers):
+      self._validate_clustered_layer(layer, clustered_layer)
+
   def testClusterModelUnsupportedKerasLayerRaisesError(self):
     """
     Verifies that attempting to cluster a model that contains an unsupported

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

@@ -57,6 +57,7 @@ def __init__(self,
                layer,
                number_of_clusters,
                cluster_centroids_init,
+               preserve_sparsity=False,
                **kwargs):
     if not isinstance(layer, Layer):
       raise ValueError(
@@ -90,9 +91,11 @@ def __init__(self,
           )
       )
 
-    if number_of_clusters <= 1:
+    limit_number_of_clusters = 2 if preserve_sparsity else 1
+    if number_of_clusters <= limit_number_of_clusters:
       raise ValueError(
-          "number_of_clusters must be greater than 1. Given: {}".format(
+          "number_of_clusters must be greater than {}. Given: {}".format(
+              limit_number_of_clusters,
               number_of_clusters
           )
       )
@@ -105,6 +108,12 @@ def __init__(self,
     # The number of cluster centroids
     self.number_of_clusters = number_of_clusters
 
+    # Whether to apply sparsity preservation or not
+    self.preserve_sparsity = preserve_sparsity
+
+    # Stores the pairs of weight names and their respective sparsity masks
+    self.sparsity_masks = {}
+
     # Stores the pairs of weight names and references to their tensors
     self.ori_weights_vars_tf = {}
 
@@ -187,7 +196,7 @@ def build(self, input_shape):
       centroid_initializer = clustering_centroids.CentroidsInitializerFactory.\
           get_centroid_initializer(
               self.cluster_centroids_init
-          )(weight, self.number_of_clusters)
+          )(weight, self.number_of_clusters, self.preserve_sparsity)
 
       cluster_centroids = centroid_initializer.get_cluster_centroids()
 
@@ -229,6 +238,16 @@ def build(self, input_shape):
           )
       )
 
+      if self.preserve_sparsity:
+        # Get the clustered weights
+        clustered_weights = self.clustering_impl[weight_name].get_clustered_weight(pulling_indices)
+
+        # Create the sparsity mask
+        sparsity_mask = tf.cast(tf.math.not_equal(clustered_weights, 0), dtype=tf.float32)
+
+        # Store the sparsity mask for training
+        self.sparsity_masks[weight_name] = sparsity_mask
+
       # We store these pairs to easily update this variables later on
       self.ori_weights_vars_tf[weight_name] = self.add_weight(
           '{}{}'.format('ori_weights_vars_tf_', weight_name),
@@ -241,13 +260,20 @@ def build(self, input_shape):
       )
 
     # We use currying here to get an updater which can be triggered at any time
-    # in future and it would return the latest version of clustered weights
+    # in the future and it would return the latest version of clustered weights
     def get_updater(for_weight_name):
       def fn():
         # Get the clustered weights
         pulling_indices = self.pulling_indices_tf[for_weight_name]
-        clustered_weights = self.clustering_impl[for_weight_name].\
-            get_clustered_weight(pulling_indices)
+        clustered_weights = self.clustering_impl[for_weight_name].get_clustered_weight(pulling_indices)
+
+        if self.preserve_sparsity:
+          # Get the sparsity mask
+          sparsity_mask = self.sparsity_masks[for_weight_name]
+
+          # Apply the sparsity mask to the clustered weights
+          clustered_weights = tf.math.multiply(clustered_weights, sparsity_mask)
+
         return clustered_weights
 
       return fn
@@ -269,19 +295,17 @@ def call(self, inputs):
     # since they are integers and not differentiable. Gradients won't flow back
     # through tf.argmin
     # Go through all tensors and replace them with their clustered copies.
-    for weight_name in self.ori_weights_vars_tf:
+    for weight_name, _ in self.clustered_vars:
+      # Get the clustered weights
       pulling_indices = self.pulling_indices_tf[weight_name]
+      clustered_weights = self.clustering_impl[weight_name].get_clustered_weight(pulling_indices)
 
-      # Update cluster associations
-      pulling_indices.assign(tf.dtypes.cast(
-          self.clustering_impl[weight_name].\
-              get_pulling_indices(self.ori_weights_vars_tf[weight_name]),
-          pulling_indices.dtype
-      ))
+      if self.preserve_sparsity:
+        # Get the sparsity mask
+        sparsity_mask = self.sparsity_masks[weight_name]
 
-      clustered_weights = self.clustering_impl[weight_name].\
-          get_clustered_weight_forward(pulling_indices,\
-              self.ori_weights_vars_tf[weight_name])
+        # Apply the sparsity mask to the clustered weights
+        clustered_weights = tf.math.multiply(clustered_weights, sparsity_mask)
 
       # Replace the weights with their clustered counterparts
       setattr(self.layer, weight_name, clustered_weights)
@@ -295,7 +319,8 @@ def get_config(self):
     base_config = super(ClusterWeights, self).get_config()
     config = {
         'number_of_clusters': self.number_of_clusters,
-        'cluster_centroids_init': self.cluster_centroids_init
+        'cluster_centroids_init': self.cluster_centroids_init,
+        'preserve_sparsity': self.preserve_sparsity
     }
     return dict(list(base_config.items()) + list(config.items()))
 
@@ -305,9 +330,11 @@ def from_config(cls, config, custom_objects=None):
 
     number_of_clusters = config.pop('number_of_clusters')
     cluster_centroids_init = config.pop('cluster_centroids_init')
+    preserve_sparsity = config.pop('preserve_sparsity')
     config['number_of_clusters'] = number_of_clusters
     config['cluster_centroids_init'] = cluster_config.CentroidInitialization(
         cluster_centroids_init)
+    config['preserve_sparsity'] = preserve_sparsity
 
     from tensorflow.python.keras.layers import deserialize as deserialize_layer  # pylint: disable=g-import-not-at-top
     layer = deserialize_layer(config.pop('layer'),

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

@@ -131,6 +131,25 @@ def testCannotBeInitializedWithNumberOfClustersLessThanTwo(
           cluster_centroids_init=CentroidInitialization.LINEAR
       )
 
+  @parameterized.parameters(
+      (0),
+      (2),
+      (-32)
+  )
+  def testCannotBeInitializedWithSparsityPreservationAndNumberOfClustersLessThanThree(
+      self, number_of_clusters):
+    """
+    Verifies that ClusterWeights cannot be initialized with less than three
+    clusters when sparsity preservation is enabled.
+    """
+    with self.assertRaises(ValueError):
+      cluster_wrapper.ClusterWeights(
+          layers.Dense(10),
+          number_of_clusters=number_of_clusters,
+          cluster_centroids_init=CentroidInitialization.LINEAR,
+          preserve_sparsity=True
+      )
+
   def testCanBeInitializedWithAlreadyClusterableLayer(self):
     """
     Verifies that ClusterWeights can be initialized with a custom clusterable