[Clustering] Enable RNN layers

[Clustering] Enable RNN layers

Change-Id: Ic44fe625f8e927bceb06a8eb19f3492b486655ef

Author: johan-gras

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

@@ -18,6 +18,7 @@
 
 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.clustering_registry import ClusteringRegistry
 
 k = tf.keras.backend
 CustomObjectScope = tf.keras.utils.CustomObjectScope
@@ -236,6 +237,14 @@ def _add_clustering_wrapper(layer):
       return layer
     if isinstance(layer, InputLayer):
       return layer.__class__.from_config(layer.get_config())
+    if isinstance(layer, tf.keras.layers.RNN):
+      return cluster_wrapper.ClusterWeightsRNN(
+        layer,
+        number_of_clusters,
+        cluster_centroids_init,
+        preserve_sparsity,
+        **kwargs,
+      )
 
     return cluster_wrapper.ClusterWeights(layer, number_of_clusters,
                                           cluster_centroids_init,
@@ -303,7 +312,7 @@ def _strip_clustering_wrapper(layer):
       for (position_variable,
            weight_name) in layer.position_original_weights.items():
         # Add the clustered weights at the correct position
-        clustered_weight = getattr(layer.layer, weight_name)
+        clustered_weight = layer.get_weight_from_layer(weight_name)
         updated_weights.insert(position_variable, clustered_weight)
 
       # Construct a clean layer with the updated weights

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

@@ -344,5 +344,152 @@ def on_train_batch_end(self, batch, logs=None):
         callbacks=[CheckWeightsCallback()])
 
 
+class ClusterRNNIntegrationTest(tf.test.TestCase, parameterized.TestCase):
+  """Integration tests for clustering RNN layers."""
+
+  def setUp(self):
+    self.max_features = 10
+    self.maxlen = 2
+    self.batch_size = 32
+    self.x_train = np.random.random((64, self.maxlen))
+    self.y_train = np.random.randint(0, 2, (64,))
+
+    self.params_clustering = {
+      "number_of_clusters": 16,
+      "cluster_centroids_init": CentroidInitialization.KMEANS_PLUS_PLUS,
+    }
+
+  def _train(self, model):
+    model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["accuracy"])
+    model.fit(
+      self.x_train,
+      self.y_train,
+      batch_size=self.batch_size,
+      epochs=1,
+    )
+
+  def _clusterTrainStrip(self, model):
+    clustered_model = cluster.cluster_weights(
+      model,
+      **self.params_clustering,
+    )
+    self._train(clustered_model)
+    stripped_model = cluster.strip_clustering(clustered_model)
+
+    return stripped_model
+
+  def __assertNbUniqueWeights(self, weight, expected_unique_weights):
+    nr_unique_weights = len(np.unique(weight.numpy().flatten()))
+    assert nr_unique_weights == expected_unique_weights
+
+  @keras_parameterized.run_all_keras_modes
+  def testClusterSimpleRNN(self):
+    model = keras.models.Sequential()
+    model.add(keras.layers.Embedding(self.max_features, 16, input_length=self.maxlen))
+    model.add(keras.layers.SimpleRNN(16, return_sequences=True))
+    model.add(keras.layers.SimpleRNN(16))
+    model.add(keras.layers.Dense(1))
+    model.add(keras.layers.Activation("sigmoid"))
+
+    stripped_model = self._clusterTrainStrip(model)
+
+    self.__assertNbUniqueWeights(
+      weight=stripped_model.layers[1].cell.kernel,
+      expected_unique_weights=self.params_clustering["number_of_clusters"],
+    )
+    self.__assertNbUniqueWeights(
+      weight=stripped_model.layers[1].cell.recurrent_kernel,
+      expected_unique_weights=self.params_clustering["number_of_clusters"],
+    )
+
+    self._train(stripped_model)
+
+  @keras_parameterized.run_all_keras_modes
+  def testClusterLSTM(self):
+    model = keras.models.Sequential()
+    model.add(keras.layers.Embedding(self.max_features, 16, input_length=self.maxlen))
+    model.add(keras.layers.LSTM(16, return_sequences=True))
+    model.add(keras.layers.LSTM(16))
+    model.add(keras.layers.Dense(1))
+    model.add(keras.layers.Activation("sigmoid"))
+
+    stripped_model = self._clusterTrainStrip(model)
+
+    self.__assertNbUniqueWeights(
+      weight=stripped_model.layers[1].cell.kernel,
+      expected_unique_weights=self.params_clustering["number_of_clusters"],
+    )
+    self.__assertNbUniqueWeights(
+      weight=stripped_model.layers[1].cell.recurrent_kernel,
+      expected_unique_weights=self.params_clustering["number_of_clusters"],
+    )
+
+    self._train(stripped_model)
+
+  @keras_parameterized.run_all_keras_modes
+  def testClusterGRU(self):
+    model = keras.models.Sequential()
+    model.add(keras.layers.Embedding(self.max_features, 16, input_length=self.maxlen))
+    model.add(keras.layers.GRU(16, return_sequences=True))
+    model.add(keras.layers.GRU(16))
+    model.add(keras.layers.Dense(1))
+    model.add(keras.layers.Activation("sigmoid"))
+
+    stripped_model = self._clusterTrainStrip(model)
+
+    self.__assertNbUniqueWeights(
+      weight=stripped_model.layers[1].cell.kernel,
+      expected_unique_weights=self.params_clustering["number_of_clusters"],
+    )
+    self.__assertNbUniqueWeights(
+      weight=stripped_model.layers[1].cell.recurrent_kernel,
+      expected_unique_weights=self.params_clustering["number_of_clusters"],
+    )
+
+    self._train(stripped_model)
+
+  @keras_parameterized.run_all_keras_modes
+  def testClusterPeepholeLSTM(self):
+    model = keras.models.Sequential()
+    model.add(keras.layers.Embedding(self.max_features, 16, input_length=self.maxlen))
+    model.add(keras.layers.RNN(tf.keras.experimental.PeepholeLSTMCell(16)))
+    model.add(keras.layers.Dense(1))
+    model.add(keras.layers.Activation("sigmoid"))
+
+    # PeepholeLSTM not supported yet.
+    with self.assertRaises(ValueError):
+      self._clusterTrainStrip(model)
+
+  @keras_parameterized.run_all_keras_modes
+  def testClusterBidirectional(self):
+    model = keras.models.Sequential()
+    model.add(keras.layers.Embedding(self.max_features, 16, input_length=self.maxlen))
+    model.add(keras.layers.Bidirectional(keras.layers.SimpleRNN(16)))
+    model.add(keras.layers.Dense(1))
+    model.add(keras.layers.Activation("sigmoid"))
+
+    # Bidirectional not supported yet.
+    with self.assertRaises(ValueError):
+      self._clusterTrainStrip(model)
+
+  @keras_parameterized.run_all_keras_modes
+  def testClusterStackedRNNCells(self):
+    model = keras.models.Sequential()
+    model.add(keras.layers.Embedding(self.max_features, 16, input_length=self.maxlen))
+    model.add(
+      tf.keras.layers.RNN(
+        tf.keras.layers.StackedRNNCells(
+          [keras.layers.SimpleRNNCell(16) for _ in range(2)]
+        )
+      )
+    )
+    model.add(keras.layers.Dense(1))
+    model.add(keras.layers.Activation("sigmoid"))
+
+    # StackedRNNCells not supported yet.
+    with self.assertRaises(ValueError):
+      self._clusterTrainStrip(model)
+
+
 if __name__ == "__main__":
   test.main()

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

@@ -139,6 +139,12 @@ def __init__(self,
   def _make_layer_name(layer):
     return '{}_{}'.format('cluster', layer.name)
 
+  def get_weight_from_layer(self, weight_name):
+    return getattr(self.layer, weight_name)
+
+  def set_weight_to_layer(self, weight_name, new_weight):
+    setattr(self.layer, weight_name, new_weight)
+
   def build(self, input_shape):
     super(ClusterWeights, self).build(input_shape)
     self.build_input_shape = input_shape
@@ -148,7 +154,7 @@ def build(self, input_shape):
       # Store the original weight in this wrapper
       # The child reference will be overridden in
       # update_clustered_weights_associations
-      original_weight = getattr(self.layer, weight_name)
+      original_weight = self.get_weight_from_layer(weight_name)
       self.original_clusterable_weights[weight_name] = original_weight
       setattr(self, 'original_weight_' + weight_name,
               original_weight)  # Track the variable
@@ -220,7 +226,7 @@ def update_clustered_weights_associations(self):
                                              self.sparsity_masks[weight_name])
 
       # Replace the weights with their clustered counterparts
-      setattr(self.layer, weight_name, clustered_weights)
+      self.set_weight_to_layer(weight_name, clustered_weights)
 
   def call(self, inputs, training=None, **kwargs):
     # Update cluster associations in order to set the latest weights
@@ -292,3 +298,13 @@ def get_weights(self):
 
   def set_weights(self, weights):
     self.layer.set_weights(weights)
+
+
+class ClusterWeightsRNN(ClusterWeights):
+  """This wrapper augments a keras RNN layer so that the weights can be clustered."""
+
+  def get_weight_from_layer(self, weight_name):
+    return getattr(self.layer.cell, weight_name)
+
+  def set_weight_to_layer(self, weight_name, new_weight):
+    setattr(self.layer.cell, weight_name, new_weight)

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

@@ -133,6 +133,18 @@ class ClusteringLookupRegistry(object):
           'kernel': ConvolutionalWeightsCA,
           'bias': BiasWeightsCA
       },
+      layers.LSTM: {
+          'kernel': DenseWeightsCA,
+          'recurrent_kernel': DenseWeightsCA
+      },
+      layers.GRU: {
+          'kernel': DenseWeightsCA,
+          'recurrent_kernel': DenseWeightsCA
+      },
+      layers.SimpleRNN: {
+          'kernel': DenseWeightsCA,
+          'recurrent_kernel': DenseWeightsCA
+      },
   }
 
   @classmethod
@@ -235,39 +247,12 @@ class ClusteringRegistry(object):
       layers.LayerNormalization: [],
   }
 
-  _RNN_CELLS_WEIGHTS_MAP = {
-      # NOTE: RNN cells are added via compat.v1 and compat.v2 to support legacy
-      # TensorFlow 2.X behavior where the v2 RNN uses the v1 RNNCell instead of
-      # the v2 RNNCell.
-      tf.compat.v1.keras.layers.GRUCell: ['kernel', 'recurrent_kernel'],
-      tf.compat.v2.keras.layers.GRUCell: ['kernel', 'recurrent_kernel'],
-      tf.compat.v1.keras.layers.LSTMCell: ['kernel', 'recurrent_kernel'],
-      tf.compat.v2.keras.layers.LSTMCell: ['kernel', 'recurrent_kernel'],
-      tf.compat.v1.keras.experimental.PeepholeLSTMCell: [
-          'kernel', 'recurrent_kernel'
-      ],
-      tf.compat.v2.keras.experimental.PeepholeLSTMCell: [
-          'kernel', 'recurrent_kernel'
-      ],
-      tf.compat.v1.keras.layers.SimpleRNNCell: ['kernel', 'recurrent_kernel'],
-      tf.compat.v2.keras.layers.SimpleRNNCell: ['kernel', 'recurrent_kernel'],
-  }
-
-  _RNN_LAYERS = {
+  _SUPPORTED_RNN_LAYERS = {
       layers.GRU,
       layers.LSTM,
-      layers.RNN,
       layers.SimpleRNN,
   }
 
-  _RNN_CELLS_STR = ', '.join(str(_RNN_CELLS_WEIGHTS_MAP.keys()))
-
-  _RNN_CELL_ERROR_MSG = (
-      'RNN Layer {} contains cell type {} which is either not supported or does'
-      'not inherit ClusterableLayer. The cell must be one of {}, or implement '
-      'ClusterableLayer.'
-  )
-
   @classmethod
   def supports(cls, layer):
     """Returns whether the registry supports this layer type.
@@ -281,31 +266,17 @@ def supports(cls, layer):
     """
     # Automatically enable layers with zero trainable weights.
     # Example: Reshape, AveragePooling2D, Maximum/Minimum, etc.
-    if not layer.trainable_weights:
+    if not layer.trainable_weights and not isinstance(layer, layers.RNN):
       return True
 
     if layer.__class__ in cls._LAYERS_WEIGHTS_MAP:
       return True
 
-    if layer.__class__ in cls._RNN_LAYERS:
-      for cell in cls._get_rnn_cells(layer):
-        if (cell.__class__ not in cls._RNN_CELLS_WEIGHTS_MAP
-            and not isinstance(cell, clusterable_layer.ClusterableLayer)):
-          return False
+    if layer.__class__ in cls._SUPPORTED_RNN_LAYERS:
       return True
 
     return False
 
-  @staticmethod
-  def _get_rnn_cells(rnn_layer):
-    if isinstance(rnn_layer.cell, layers.StackedRNNCells):
-      return rnn_layer.cell.cells
-    return [rnn_layer.cell]
-
-  @classmethod
-  def _is_rnn_layer(cls, layer):
-    return layer.__class__ in cls._RNN_LAYERS
-
   @classmethod
   def _weight_names(cls, layer):
     # For layers with zero trainable weights, like Reshape, Pooling.
@@ -323,7 +294,6 @@ def make_clusterable(cls, layer):
 
     Returns:
       The modified layer object.
-
     """
 
     if not cls.supports(layer):
@@ -334,23 +304,16 @@ def get_clusterable_weights():
               for weight_name in cls._weight_names(layer)]
 
     def get_clusterable_weights_rnn():  # pylint: disable=missing-docstring
-      def get_clusterable_weights_rnn_cell(cell):
-        if cell.__class__ in cls._RNN_CELLS_WEIGHTS_MAP:
-          return [(weight, getattr(cell, weight))
-                  for weight in cls._RNN_CELLS_WEIGHTS_MAP[cell.__class__]]
-
-        if isinstance(cell, clusterable_layer.ClusterableLayer):
-          return cell.get_clusterable_weights()
-
-        raise ValueError(cls._RNN_CELL_ERROR_MSG.format(
-            layer.__class__, cell.__class__, cls._RNN_CELLS_WEIGHTS_MAP.keys()))
+      if isinstance(layer.cell, clusterable_layer.ClusterableLayer):
+        raise ValueError("ClusterableLayer is not yet supported for RNNs based layer.")
 
-      clusterable_weights = []
-      for rnn_cell in cls._get_rnn_cells(layer):
-        clusterable_weights.extend(get_clusterable_weights_rnn_cell(rnn_cell))
+      clusterable_weights = [
+        ('kernel', layer.cell.kernel),
+        ('recurrent_kernel', layer.cell.recurrent_kernel),
+      ]
       return clusterable_weights
 
-    if cls._is_rnn_layer(layer):
+    if layer.__class__ in cls._SUPPORTED_RNN_LAYERS:
       layer.get_clusterable_weights = get_clusterable_weights_rnn
     else:
       layer.get_clusterable_weights = get_clusterable_weights