[Clustering] Average Gradient Aggregation

* Add Average Gradient Aggregation parameter to the clustering API
* Cleaning of the wrapper logic

Change-Id: Ia7f94a1bc8de2e13663ddcf676748413ff1f9295

Author: johan-gras

tensorflow_model_optimization/python/core/clustering/keras/BUILD

@@ -82,6 +82,7 @@ py_library(
     visibility = ["//visibility:public"],
     deps = [
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config",
     ],
 )
 

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

@@ -293,36 +293,25 @@ def _strip_clustering_wrapper(layer):
     if isinstance(layer, keras.Model):
       return keras.models.clone_model(
           layer, input_tensors=None, clone_function=_strip_clustering_wrapper)
+
     elif isinstance(layer, cluster_wrapper.ClusterWeights):
-      if not hasattr(layer.layer, '_batch_input_shape') and\
-          hasattr(layer, '_batch_input_shape'):
-        # pylint: disable=protected-access
-        layer.layer._batch_input_shape = layer._batch_input_shape
+      # Update cluster associations in order to get the latest weights
+      layer.update_clustered_weights_associations()
+
+      # Construct a list of weights to initialize the clean layer
+      updated_weights = layer.layer.get_weights()  # non clusterable weights only
+      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)
+        updated_weights.insert(position_variable, clustered_weight)
+
+      # Construct a clean layer with the updated weights
+      clean_layer = layer.layer.from_config(layer.layer.get_config())
+      clean_layer.build(layer.build_input_shape)
+      clean_layer.set_weights(updated_weights)
+
+      return clean_layer
 
-      # We reset both arrays of weights, so that we can guarantee the correct
-      # order of newly created weights
-      # pylint: disable=protected-access
-      layer.layer._trainable_weights = []
-      layer.layer._non_trainable_weights = []
-      for i in range(len(layer.restore)):
-        # This is why we used integers as keys
-        name, weight_name, weight = layer.restore[i]
-        # In both cases we use k.batch_get_value since we need physical copies
-        # of the arrays to initialize a new tensor
-        if i in layer.gone_variables:
-          # If the variable was removed because it was clustered, we restore it
-          # by using updater we created earlier
-          new_weight_value = k.batch_get_value([weight()])[0]
-        else:
-          # If the value was not clustered(e.g. bias), we still store a valid
-          # reference to the tensor. We use this reference to get the value
-          new_weight_value = k.batch_get_value([weight])[0]
-        setattr(layer.layer,
-                name,
-                k.variable(new_weight_value, name=weight_name))
-      # When all weights are filled with the values, just return the underlying
-      # layer since it is now fully autonomous from its wrapper
-      return layer.layer
     return layer
 
   # Just copy the model with the right callback

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

@@ -31,7 +31,19 @@ class CentroidInitialization(str, enum.Enum):
        initialize the clusters centroids.
   * `KMEANS_PLUS_PLUS`: cluster centroids using the kmeans++ algorithm
   """
-  LINEAR = "LINEAR"
-  RANDOM = "RANDOM"
-  DENSITY_BASED = "DENSITY_BASED"
-  KMEANS_PLUS_PLUS = "KMEANS_PLUS_PLUS"
+  LINEAR = "CentroidInitialization.LINEAR"
+  RANDOM = "CentroidInitialization.RANDOM"
+  DENSITY_BASED = "CentroidInitialization.DENSITY_BASED"
+  KMEANS_PLUS_PLUS = "CentroidInitialization.KMEANS_PLUS_PLUS"
+
+
+class GradientAggregation(str, enum.Enum):
+  """Specifies how the cluster gradient should be aggregated.
+
+  * `SUM`: The gradient of each cluster centroid is the sum of their
+      respective child’s weight gradient.
+  * `AVG`: The gradient of each cluster centroid is the averaged sum of
+      their respective child’s weight gradient.
+  """
+  SUM = "GradientAggregation.SUM"
+  AVG = "GradientAggregation.AVG"

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

@@ -43,7 +43,6 @@ def setUp(self):
         "cluster_centroids_init": CentroidInitialization.LINEAR
     }
 
-
   @parameterized.parameters(_distribution_strategies())
   def testClusterSimpleDenseModel(self, distribution):
     """End-to-end test."""
@@ -64,7 +63,7 @@ def testClusterSimpleDenseModel(self, distribution):
     model.predict(np.random.rand(20, 10))
 
     stripped_model = cluster.strip_clustering(model)
-    weights_as_list = stripped_model.get_weights()[0].reshape(-1,).tolist()
+    weights_as_list = stripped_model.layers[0].kernel.numpy().reshape(-1,).tolist()
     unique_weights = set(weights_as_list)
     self.assertLessEqual(len(unique_weights), self.params["number_of_clusters"])
 
@@ -87,7 +86,7 @@ def testAssociationValuesPerReplica(self, distribution):
       self.assertEqual(len(clusterable_weights), 1)
       weights_name = clusterable_weights[0][0]
       self.assertEqual(weights_name, 'kernel')
-      centroids1 = l.cluster_centroids_tf[weights_name]
+      centroids1 = l.cluster_centroids[weights_name]
 
       mean_weight = tf.reduce_mean(l.layer.kernel)
       min_weight = tf.reduce_min(l.layer.kernel)
@@ -119,18 +118,18 @@ def update_fn(v, val):
           centroids1, update_fn, args=(initial_val,))
       l.call(tf.ones(shape=input_shape))
 
-      clst_indices = l.pulling_indices_tf[weights_name]
+      clst_indices = l.pulling_indices[weights_name]
       per_replica = distribution.experimental_local_results(clst_indices)
       assert_all_cluster_indices(per_replica, 0)
 
       second_val = tf.Variable([mean_weight - 2.0 * max_dist, mean_weight], \
         aggregation=tf.VariableAggregation.MEAN)
-      centroids2 = l.cluster_centroids_tf[weights_name]
+      centroids2 = l.cluster_centroids[weights_name]
       centroids2 = distribution.extended.update(
           centroids2, update_fn, args=(second_val,))
       l.call(tf.ones(shape=input_shape))
 
-      clst_indices = l.pulling_indices_tf[weights_name]
+      clst_indices = l.pulling_indices[weights_name]
       per_replica = distribution.experimental_local_results(clst_indices)
       assert_all_cluster_indices(per_replica, 1)
 

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

@@ -174,7 +174,7 @@ def testSparsityIsPreservedDuringTraining(self):
         original_model, **clustering_params)
 
     stripped_model_before_tuning = cluster.strip_clustering(clustered_model)
-    weights_before_tuning = stripped_model_before_tuning.get_weights()[0]
+    weights_before_tuning = stripped_model_before_tuning.layers[0].kernel
     non_zero_weight_indices_before_tuning = np.nonzero(weights_before_tuning)
 
     clustered_model.compile(
@@ -185,9 +185,9 @@ def testSparsityIsPreservedDuringTraining(self):
     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]
+    weights_after_tuning = stripped_model_after_tuning.layers[0].kernel
     non_zero_weight_indices_after_tuning = np.nonzero(weights_after_tuning)
-    weights_as_list_after_tuning = weights_after_tuning.reshape(-1,).tolist()
+    weights_as_list_after_tuning = weights_after_tuning.numpy().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.
@@ -299,6 +299,54 @@ def clusters_check(stripped_model):
 
     self.end_to_end_testing(original_model, clusters_check)
 
+  @keras_parameterized.run_all_keras_modes
+  def testWeightsAreLearningDuringClustering(self):
+    """Verifies that training a clustered model does update
+    original_weights, clustered_centroids and bias."""
+    original_model = keras.Sequential([
+      layers.Dense(5, input_shape=(5,))
+    ])
+
+    clustered_model = cluster.cluster_weights(original_model, **self.params)
+
+    clustered_model.compile(
+      loss=keras.losses.categorical_crossentropy,
+      optimizer="adam",
+      metrics=["accuracy"],
+    )
+
+    class CheckWeightsCallback(keras.callbacks.Callback):
+      def on_train_batch_begin(self, batch, logs=None):
+        # Save weights before batch
+        self.original_weight_kernel = (
+          self.model.layers[0].original_clusterable_weights['kernel'].numpy()
+        )
+        self.cluster_centroids_kernel = (
+          self.model.layers[0].cluster_centroids['kernel'].numpy()
+        )
+        self.bias = (
+          self.model.layers[0].layer.bias.numpy()
+        )
+
+      def on_train_batch_end(self, batch, logs=None):
+        # Check weights are different after batch
+        assert not np.array_equal(
+          self.original_weight_kernel,
+          self.model.layers[0].original_clusterable_weights['kernel'].numpy()
+        )
+        assert not np.array_equal(
+          self.cluster_centroids_kernel,
+          self.model.layers[0].cluster_centroids['kernel'].numpy()
+        )
+        assert not np.array_equal(
+          self.bias,
+          self.model.layers[0].layer.bias.numpy()
+        )
+
+    clustered_model.fit(x=self.dataset_generator(),
+                        steps_per_epoch=5,
+                        callbacks=[CheckWeightsCallback()])
+
 
 if __name__ == "__main__":
   test.main()

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

@@ -16,6 +16,9 @@
 
 import json
 
+import tempfile
+
+import os
 from absl.testing import parameterized
 import tensorflow as tf
 
@@ -253,6 +256,76 @@ def testClusterKerasCustomLayer(self):
     with self.assertRaises(ValueError):
       cluster_wrapper.ClusterWeights(keras_custom_layer, **self.params)
 
+  def testStripClusteringSequentialModelWithKernelRegularizer(self):
+    """Verifies that stripping the clustering wrappers from a sequential model produces the expected config."""
+    model = keras.Sequential([
+        layers.Dense(10, input_shape=(10,)),
+        layers.Dense(10, kernel_regularizer=tf.keras.regularizers.L1(0.01)),
+    ])
+    clustered_model = cluster.cluster_weights(model, **self.params)
+    stripped_model = cluster.strip_clustering(clustered_model)
+    # check that kernel regularizer is present in the second dense layer
+    self.assertIsNotNone(stripped_model.layers[1].kernel_regularizer)
+    with tempfile.TemporaryDirectory() as tmp_dir_name:
+      keras_file = os.path.join(tmp_dir_name, 'cluster_test')
+      stripped_model.save(keras_file, save_traces = True)
+
+  def testStripClusteringSequentialModelWithBiasRegularizer(self):
+    """Verifies that stripping the clustering wrappers from a sequential model produces the expected config."""
+    model = keras.Sequential([
+        layers.Dense(10, input_shape=(10,)),
+        layers.Dense(10, bias_regularizer=tf.keras.regularizers.L1(0.01)),
+    ])
+    clustered_model = cluster.cluster_weights(model, **self.params)
+    stripped_model = cluster.strip_clustering(clustered_model)
+    # check that kernel regularizer is present in the second dense layer
+    self.assertIsNotNone(stripped_model.layers[1].bias_regularizer)
+    with tempfile.TemporaryDirectory() as tmp_dir_name:
+      keras_file = os.path.join(tmp_dir_name, 'cluster_test')
+      stripped_model.save(keras_file, save_traces = True)
+
+  def testStripClusteringSequentialModelWithActivityRegularizer(self):
+    """Verifies that stripping the clustering wrappers from a sequential model produces the expected config."""
+    model = keras.Sequential([
+        layers.Dense(10, input_shape=(10,)),
+        layers.Dense(10, activity_regularizer=tf.keras.regularizers.L1(0.01)),
+    ])
+    clustered_model = cluster.cluster_weights(model, **self.params)
+    stripped_model = cluster.strip_clustering(clustered_model)
+    # check that kernel regularizer is present in the second dense layer
+    self.assertIsNotNone(stripped_model.layers[1].activity_regularizer)
+    with tempfile.TemporaryDirectory() as tmp_dir_name:
+      keras_file = os.path.join(tmp_dir_name, 'cluster_test')
+      stripped_model.save(keras_file, save_traces = True)
+
+  def testStripClusteringSequentialModelWithKernelConstraint(self):
+    """Verifies that stripping the clustering wrappers from a sequential model produces the expected config."""
+    model = keras.Sequential([
+        layers.Dense(10, input_shape=(10,)),
+        layers.Dense(10, kernel_constraint=tf.keras.constraints.max_norm(2.)),
+    ])
+    clustered_model = cluster.cluster_weights(model, **self.params)
+    stripped_model = cluster.strip_clustering(clustered_model)
+    # check that kernel regularizer is present in the second dense layer
+    self.assertIsNotNone(stripped_model.layers[1].kernel_constraint)
+    with tempfile.TemporaryDirectory() as tmp_dir_name:
+      keras_file = os.path.join(tmp_dir_name, 'cluster_test')
+      stripped_model.save(keras_file, save_traces = True)
+
+  def testStripClusteringSequentialModelWithBiasConstraint(self):
+    """Verifies that stripping the clustering wrappers from a sequential model produces the expected config."""
+    model = keras.Sequential([
+        layers.Dense(10, input_shape=(10,)),
+        layers.Dense(10, bias_constraint=tf.keras.constraints.max_norm(2.)),
+    ])
+    clustered_model = cluster.cluster_weights(model, **self.params)
+    stripped_model = cluster.strip_clustering(clustered_model)
+    # check that kernel regularizer is present in the second dense layer
+    self.assertIsNotNone(stripped_model.layers[1].bias_constraint)
+    with tempfile.TemporaryDirectory() as tmp_dir_name:
+      keras_file = os.path.join(tmp_dir_name, 'cluster_test')
+      stripped_model.save(keras_file, save_traces = True)
+
   def testClusterMyClusterableLayer(self):
     # we have weights to cluster.
     layer = self.clusterable_layer
@@ -539,7 +612,7 @@ def testClusterSubclassModelAsSubmodel(self):
   def testStripClusteringSequentialModel(self):
     """Verifies that stripping the clustering wrappers from a sequential model produces the expected config."""
     model = keras.Sequential([
-        layers.Dense(10),
+        layers.Dense(10, input_shape=(5,)),
         layers.Dense(10),
     ])
 
@@ -582,7 +655,7 @@ def testClusterWeightsStrippedWeights(self):
 
   @keras_parameterized.run_all_keras_modes
   def testStrippedKernel(self):
-    """Verifies that stripping the clustering wrappers from a functional model restores the layers kernel and the layers weight array to the new clustered weight value ."""
+    """Verifies that stripping the clustering wrappers from a functional model restores the layers kernel and the layers weight array to the new clustered weight value."""
     i1 = keras.Input(shape=(1, 1, 1))
     x1 = layers.Conv2D(1, 1)(i1)
     outputs = x1
@@ -596,8 +669,7 @@ def testStrippedKernel(self):
 
     self.assertEqual(self._count_clustered_layers(stripped_model), 0)
     self.assertIsNot(stripped_conv2d_layer.kernel, clustered_kernel)
-    self.assertEqual(stripped_conv2d_layer.kernel,
-                     stripped_conv2d_layer.weights[0])
+    self.assertIn(stripped_conv2d_layer.kernel, stripped_conv2d_layer.weights)
 
   @keras_parameterized.run_all_keras_modes
   def testStripSelectivelyClusteredFunctionalModel(self):
@@ -628,5 +700,23 @@ def testStripSelectivelyClusteredSequentialModel(self):
     self.assertEqual(self._count_clustered_layers(stripped_model), 0)
     self.assertIsInstance(stripped_model.layers[0], layers.Dense)
 
+  @keras_parameterized.run_all_keras_modes
+  def testStripClusteringAndSetOriginalWeightsBack(self):
+    """Verifies that we can set_weights onto the stripped model."""
+    model = keras.Sequential([
+      layers.Dense(10, input_shape=(5,)),
+      layers.Dense(10),
+    ])
+
+    # Save original weights
+    original_weights = model.get_weights()
+
+    # Cluster and strip
+    clustered_model = cluster.cluster_weights(model, **self.params)
+    stripped_model = cluster.strip_clustering(clustered_model)
+
+    # Set back original weights onto the strip model
+    stripped_model.set_weights(original_weights)
+
 if __name__ == '__main__':
   test.main()