Support of MHA for clustering.

Change-Id: I7005bffa70770ce6e36859384dfca351e39318aa

Author: wwwind

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

@@ -244,6 +244,14 @@ def _add_clustering_wrapper(layer):
           preserve_sparsity,
           **kwargs,
       )
+    if isinstance(layer, tf.keras.layers.MultiHeadAttention):
+      return cluster_wrapper.ClusterWeightsMHA(
+          layer,
+          number_of_clusters,
+          cluster_centroids_init,
+          preserve_sparsity,
+          **kwargs,
+      )
 
     return cluster_wrapper.ClusterWeights(layer, number_of_clusters,
                                           cluster_centroids_init,
@@ -301,6 +309,10 @@ def _strip_clustering_wrapper(layer):
       return tf.keras.models.clone_model(
           layer, input_tensors=None, clone_function=_strip_clustering_wrapper)
 
+    elif isinstance(layer, cluster_wrapper.ClusterWeightsMHA):
+        # In case of MHA layer, use the overloaded implementation
+        return layer.strip_clustering()
+
     elif isinstance(layer, cluster_wrapper.ClusterWeights):
       # Update cluster associations in order to get the latest weights
       layer.update_clustered_weights_associations()

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

@@ -534,5 +534,47 @@ def testClusterStackedRNNCells(self):
         expected_unique_weights=self.params_clustering["number_of_clusters"],
     )
 
+class ClusterMHAIntegrationTest(tf.test.TestCase, parameterized.TestCase):
+  """Integration tests for clustering MHA layer."""
+
+  def setUp(self):
+    self.x_train = np.random.uniform(size=(500, 32, 32))
+    self.y_train = np.random.randint(low=0, high=1024, size=(500,))
+
+    self.nr_of_clusters = 16
+    self.params_clustering = {
+      "number_of_clusters": self.nr_of_clusters,
+      "cluster_centroids_init": CentroidInitialization.KMEANS_PLUS_PLUS,
+    }
+
+  def _get_model(self):
+    """Returns functional model with MHA layer."""
+    inp = tf.keras.layers.Input(shape=(32,32), batch_size=100)
+    x = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=16)(query=inp, value=inp)
+    out = tf.keras.layers.Flatten()(x)
+    model = tf.keras.Model(inputs=inp, outputs=out)
+    return model
+
+  @keras_parameterized.run_all_keras_modes
+  def testMHA(self):
+    model = self._get_model()
+
+    clustered_model = cluster.cluster_weights(model, **self.params_clustering)
+
+    clustered_model.compile(
+      optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),
+      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+      metrics=[tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy')])
+    clustered_model.run_eagerly = True
+    clustered_model.fit(self.x_train, self.y_train, epochs=1, batch_size=100, verbose=1)
+
+    stripped_model = cluster.strip_clustering(clustered_model)
+
+    layerMHA = stripped_model.layers[1]
+    for weight in layerMHA.weights:
+      if 'kernel' in weight.name:
+        nr_unique_weights = len(np.unique(weight.numpy()))
+        assert nr_unique_weights == self.nr_of_clusters
+
 if __name__ == "__main__":
   test.main()

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

@@ -14,6 +14,8 @@
 # ==============================================================================
 """Keras ClusterWeights wrapper API."""
 
+from operator import attrgetter
+
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
@@ -388,3 +390,32 @@ def set_weight_to_layer(self, weight_name, new_weight):
       return setattr(return_layer_cell, weight_name_no_index, new_weight)
     else:
       raise ValueError('No cells in the RNN layer to set weights for.')
+
+
+class ClusterWeightsMHA(ClusterWeights):
+  """This wrapper augments a keras MHA layer so that the weights can be clustered."""
+
+  def get_weight_from_layer(self, weight_name):
+    pre, _, post = weight_name.rpartition('.')
+    return getattr(getattr(self.layer, pre), post)
+
+  def set_weight_to_layer(self, weight_name, new_weight):
+    pre, _, post = weight_name.rpartition('.')
+    layer = attrgetter(pre)(self.layer)
+    setattr(layer, post, new_weight)
+
+  def strip_clustering(self):
+    """ The restore from config is not working for MHA layer, because
+    weights are not created when the build function is called. Therefore,
+    original weights have been replaced in the layer."""
+    for weight_name, original_weight in self.original_clusterable_weights.items():
+
+      # Get the clustered weights
+      clustered_weight = self.get_weight_from_layer(weight_name)
+
+      # Re-assign these weights to the original
+      original_weight.assign(clustered_weight)
+      setattr(self.layer, weight_name, original_weight)
+
+    return self.layer
+

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

@@ -92,6 +92,10 @@ class ClusteringRegistry(object):
       layers.Bidirectional,
   ])
 
+  _SUPPORTED_MHA_LAYERS = {
+      tf.keras.layers.MultiHeadAttention,
+  }
+
   @classmethod
   def supports(cls, layer):
     """Returns whether the registry supports this layer type.
@@ -121,6 +125,9 @@ def supports(cls, layer):
           return False
       return True
 
+    if layer.__class__ in cls._SUPPORTED_MHA_LAYERS:
+      return True
+
     return False
 
   def _get_rnn_cells(rnn_layer):  # pylint: disable=no-self-argument
@@ -190,8 +197,16 @@ def get_clusterable_weights_rnn_cell(cell, i):
           clusterable_weights = get_clusterable_weights_rnn_cell(rnn_cell, 0)
       return clusterable_weights
 
+    def get_clusterable_weights_mha():  # pylint: disable=missing-docstring
+      return [('_query_dense.kernel', layer._query_dense.kernel),
+              ('_key_dense.kernel', layer._key_dense.kernel),
+              ('_value_dense.kernel', layer._value_dense.kernel),
+              ('_output_dense.kernel', layer._output_dense.kernel)]
+
     if layer.__class__ in cls._SUPPORTED_RNN_LAYERS:
       layer.get_clusterable_weights = get_clusterable_weights_rnn
+    elif layer.__class__ in cls._SUPPORTED_MHA_LAYERS:
+      layer.get_clusterable_weights = get_clusterable_weights_mha
     else:
       layer.get_clusterable_weights = get_clusterable_weights
 

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

@@ -512,6 +512,13 @@ def testMakeClusterableRaisesErrorOnRNNLayersUnsupportedCell(self):
     with self.assertRaises(ValueError):
       ClusterRegistry.make_clusterable(layer)
 
+  def testSupportsMultiHeadAttentionLayer(self):
+    """
+    Verifies that ClusterRegistry supports a MultiHeadAttention layer.
+    """
+    layer = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=2)
+    self.assertTrue(ClusterRegistry.supports(layer))
+    ClusterRegistry.make_clusterable(layer)
 
 if __name__ == '__main__':
   test.main()