Merge pull request #520 from MatteoArm:feature/sparsity_aware_clustering

PiperOrigin-RevId: 351486381

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/api/clustering/keras/__init__.py

@@ -14,6 +14,8 @@
 # ==============================================================================
 """Module containing clustering code built on Keras abstractions."""
 # pylint: disable=g-bad-import-order
+from tensorflow_model_optimization.python.core.clustering.keras import experimental
+
 from tensorflow_model_optimization.python.core.clustering.keras.cluster import cluster_scope
 from tensorflow_model_optimization.python.core.clustering.keras.cluster import cluster_weights
 from tensorflow_model_optimization.python.core.clustering.keras.cluster import strip_clustering

tensorflow_model_optimization/python/core/api/clustering/keras/experimental/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+# ==============================================================================
+"""Module containing experimental clustering code built on Keras abstractions."""
+from tensorflow_model_optimization.python.core.clustering.keras.experimental.cluster import cluster_weights

tensorflow_model_optimization/python/core/clustering/keras/BUILD

@@ -2,7 +2,7 @@ package(default_visibility = [
     "//tensorflow_model_optimization:__subpackages__",
 ])
 
-licenses(["notice"])  # Apache 2.0
+licenses(["notice"])
 
 py_library(
     name = "keras",
@@ -12,6 +12,7 @@ py_library(
     srcs_version = "PY3",
     deps = [
         ":cluster",
+        "//tensorflow_model_optimization/python/core/clustering/keras/experimental",
     ],
 )
 
@@ -92,6 +93,7 @@ py_test(
     deps = [
         ":cluster",
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/clustering/keras/experimental:cluster",
     ],
 )
 
@@ -146,6 +148,7 @@ py_test(
     deps = [
         ":cluster",
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/clustering/keras/experimental:cluster",
         "//tensorflow_model_optimization/python/core/keras:compat",
     ],
 )

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

@@ -15,9 +15,7 @@
 """Clustering API functions for Keras models."""
 
 from tensorflow import keras
-from tensorflow.keras import initializers
 
-from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_wrapper
 from tensorflow_model_optimization.python.core.clustering.keras import clustering_centroids
 
@@ -57,7 +55,7 @@ def cluster_weights(to_cluster,
                     number_of_clusters,
                     cluster_centroids_init,
                     **kwargs):
-  """Modify a keras layer or model to be clustered during training.
+  """Modifies a keras layer or model to be clustered during training.
 
   This function wraps a keras model or layer with clustering functionality
   which clusters the layer's weights during training. For examples, using
@@ -80,8 +78,7 @@ def cluster_weights(to_cluster,
   ```python
   clustering_params = {
     'number_of_clusters': 8,
-    'cluster_centroids_init':
-        CentroidInitialization.DENSITY_BASED
+    'cluster_centroids_init': CentroidInitialization.DENSITY_BASED
   }
 
   clustered_model = cluster_weights(original_model, **clustering_params)
@@ -92,8 +89,104 @@ def cluster_weights(to_cluster,
   ```python
   clustering_params = {
     'number_of_clusters': 8,
-    'cluster_centroids_init':
-        CentroidInitialization.DENSITY_BASED
+    'cluster_centroids_init': CentroidInitialization.DENSITY_BASED
+  }
+
+  model = keras.Sequential([
+      layers.Dense(10, activation='relu', input_shape=(100,)),
+      cluster_weights(layers.Dense(2, activation='tanh'), **clustering_params)
+  ])
+  ```
+
+  Arguments:
+      to_cluster: A single keras layer, list of keras layers, or a
+        `tf.keras.Model` instance.
+      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: 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
+      **kwargs: Additional keyword arguments to be passed to the keras layer.
+        Ignored when to_cluster is not a keras layer.
+
+  Returns:
+    Layer or model modified to include clustering related metadata.
+
+  Raises:
+    ValueError: if the keras layer is unsupported, or the keras model contains
+    an unsupported layer.
+  """
+  return _cluster_weights(
+      to_cluster,
+      number_of_clusters,
+      cluster_centroids_init,
+      preserve_sparsity=False,
+      **kwargs)
+
+
+def _cluster_weights(to_cluster, number_of_clusters, cluster_centroids_init,
+                     preserve_sparsity, **kwargs):
+  """Modifies a keras layer or model to be clustered during training.
+
+  This function wraps a keras model or layer with clustering functionality
+  which clusters the layer's weights during training. For examples, using
+  this with number_of_clusters equals 8 will ensure that each weight tensor has
+  no more than 8 unique values.
+
+  Before passing to the clustering API, a model should already be trained and
+  show some acceptable performance on the testing/validation sets.
+
+  The function accepts either a single keras layer
+  (subclass of `keras.layers.Layer`), list of keras layers or a keras model
+  (instance of `keras.models.Model`) and handles them appropriately.
+
+  If it encounters a layer it does not know how to handle, it will throw an
+  error. While clustering an entire model, even a single unknown layer would
+  lead to an error.
+
+  Cluster a model:
+
+  ```python
+  clustering_params = {
+    'number_of_clusters': 8,
+    'cluster_centroids_init': CentroidInitialization.DENSITY_BASED,
+    'preserve_sparsity': False
+  }
+
+  clustered_model = cluster_weights(original_model, **clustering_params)
+  ```
+
+  Cluster a layer:
+
+  ```python
+  clustering_params = {
+    'number_of_clusters': 8,
+    'cluster_centroids_init': CentroidInitialization.DENSITY_BASED,
+    'preserve_sparsity': False
+  }
+
+  model = keras.Sequential([
+      layers.Dense(10, activation='relu', input_shape=(100,)),
+      cluster_weights(layers.Dense(2, activation='tanh'), **clustering_params)
+  ])
+  ```
+
+  Cluster a layer with sparsity preservation (experimental):
+
+  ```python
+  clustering_params = {
+    'number_of_clusters': 8,
+    'cluster_centroids_init': CentroidInitialization.DENSITY_BASED,
+    'preserve_sparsity': True
   }
 
   model = keras.Sequential([
@@ -110,6 +203,8 @@ def cluster_weights(to_cluster,
         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.
+      preserve_sparsity (experimental): 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.
 
@@ -120,33 +215,31 @@ def cluster_weights(to_cluster,
     ValueError: if the keras layer is unsupported, or the keras model contains
     an unsupported layer.
   """
-  if not clustering_centroids.CentroidsInitializerFactory.\
-      init_is_supported(cluster_centroids_init):
-    raise ValueError("Cluster centroid initialization {} not supported".\
-        format(cluster_centroids_init))
+  if not clustering_centroids.CentroidsInitializerFactory.init_is_supported(
+      cluster_centroids_init):
+    raise ValueError('Cluster centroid initialization {} not supported'.format(
+        cluster_centroids_init))
 
   def _add_clustering_wrapper(layer):
-
-    if (isinstance(layer, keras.Model)):
+    if isinstance(layer, keras.Model):
       # Check whether the model is a subclass.
       # NB: This check is copied from keras.py file in tensorflow.
       # There is no available public API to do this check.
+      # pylint: disable=protected-access
       if (not layer._is_graph_network and
           not isinstance(layer, keras.models.Sequential)):
-        raise ValueError("Subclassed models are not supported currently.")
+        raise ValueError('Subclassed models are not supported currently.')
 
-      return keras.models.clone_model(layer,
-                                    input_tensors=None,
-                                    clone_function=_add_clustering_wrapper)
+      return keras.models.clone_model(
+          layer, input_tensors=None, clone_function=_add_clustering_wrapper)
     if isinstance(layer, cluster_wrapper.ClusterWeights):
       return layer
     if isinstance(layer, InputLayer):
       return layer.__class__.from_config(layer.get_config())
 
-    return cluster_wrapper.ClusterWeights(layer,
-                                          number_of_clusters,
+    return cluster_wrapper.ClusterWeights(layer, number_of_clusters,
                                           cluster_centroids_init,
-                                          **kwargs)
+                                          preserve_sparsity, **kwargs)
 
   def _wrap_list(layers):
     output = []
@@ -166,7 +259,7 @@ def _wrap_list(layers):
 
 
 def strip_clustering(model):
-  """Strip clustering wrappers from the model.
+  """Strips clustering wrappers from the model.
 
   Once a model has been clustered, this method can be used
   to restore the original model with the clustered weights.
@@ -198,16 +291,17 @@ def strip_clustering(model):
 
   def _strip_clustering_wrapper(layer):
     if isinstance(layer, keras.Model):
-      return keras.models.clone_model(layer,
-                               input_tensors=None,
-                               clone_function=_strip_clustering_wrapper)
+      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
 
       # 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)):