Merge pull request #845 from wwwind:mha_clustering_support

PiperOrigin-RevId: 402726607

Author: tensorflower-gardener

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,13 @@ 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):
+        # Update cluster associations in order to get the latest weights
+        layer.update_clustered_weights_associations()
+
+        # 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,46 @@ 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.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()

tensorflow_model_optimization/python/examples/clustering/keras/mnist/mnist_mha.py

@@ -0,0 +1,106 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# pylint: disable=missing-docstring
+"""Train a simple convnet with MultiHeadAttention layer on MNIST dataset
+and cluster it.
+"""
+import tensorflow as tf
+import tensorflow_model_optimization as tfmot
+
+import numpy as np
+
+NUMBER_OF_CLUSTERS = 3
+
+# Load MNIST dataset
+mnist = tf.keras.datasets.mnist
+(train_images, train_labels), (test_images, test_labels) = mnist.load_data()
+
+# Normalize the input image so that each pixel value is between 0 to 1.
+train_images = train_images / 255.0
+test_images = test_images / 255.0
+
+# define model
+input = tf.keras.layers.Input(shape=(28, 28))
+x = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=16, name="mha")(
+    query=input, value=input
+)
+x = tf.keras.layers.Flatten()(x)
+out = tf.keras.layers.Dense(10)(x)
+model = tf.keras.Model(inputs=input, outputs=out)
+
+# Train the digit classification model
+model.compile(
+    optimizer="adam",
+    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+    metrics=["accuracy"],
+)
+
+model.fit(
+    train_images, train_labels, epochs=1, validation_split=0.1,
+)
+
+score = model.evaluate(test_images, test_labels, verbose=0)
+print('Model test loss:', score[0])
+print('Model test accuracy:', score[1])
+
+# Compute end step to finish pruning after 2 epochs.
+batch_size = 128
+epochs = 1
+validation_split = 0.1  # 10% of training set will be used for validation set.
+
+# Define model for clustering
+cluster_weights = tfmot.clustering.keras.cluster_weights
+CentroidInitialization = tfmot.clustering.keras.CentroidInitialization
+
+clustering_params = {
+    "number_of_clusters": NUMBER_OF_CLUSTERS,
+    "cluster_centroids_init": CentroidInitialization.KMEANS_PLUS_PLUS,
+}
+model_for_clustering = cluster_weights(model, **clustering_params)
+
+# `cluster_weights` requires a recompile.
+model_for_clustering.compile(
+    optimizer="adam",
+    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+    metrics=["accuracy"],
+)
+
+model_for_clustering.fit(
+    train_images,
+    train_labels,
+    batch_size=batch_size,
+    epochs=epochs,
+    validation_split=validation_split,
+)
+
+score = model_for_clustering.evaluate(test_images, test_labels, verbose=0)
+print('Clustered model test loss:', score[0])
+print('Clustered model test accuracy:', score[1])
+
+# Strip clustering from the model
+clustered_model = tfmot.clustering.keras.strip_clustering(model_for_clustering)
+clustered_model.compile(
+    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+    optimizer='adam',
+    metrics=['accuracy'])
+
+score = clustered_model.evaluate(test_images, test_labels, verbose=0)
+print('Stripped clustered model test loss:', score[0])
+print('Stripped clustered model test accuracy:', score[1])
+
+# Check that numbers of weights for MHA layer is the given number of clusters.
+mha_weights = list(filter(lambda x: 'mha' in x.name and 'kernel' in x.name, clustered_model.weights))
+for x in mha_weights:
+    assert len(np.unique(x.numpy())) == NUMBER_OF_CLUSTERS