Add Anti Zero-Drift functionality for Sparsity-Aware clustering (experimental)

 * Created new experimental API for sparsity-aware clustering
 * Kept the original API implementation
 * Moved the new feature to a new experimental package, making the
   original implementation private
 * Updated the unit tests accordingly
 * Created init and BUILD files for the new experimental package

Author: MatteoArm

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 2020 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

@@ -12,6 +12,7 @@ py_library(
     srcs_version = "PY3",
     deps = [
         ":cluster",
+        "//tensorflow_model_optimization/python/core/clustering/keras/experimental",
     ],
 )
 
@@ -90,6 +91,7 @@ py_test(
     visibility = ["//visibility:public"],
     deps = [
         ":cluster",
+        "//tensorflow_model_optimization/python/core/clustering/keras/experimental:cluster",
         # tensorflow dep1,
     ],
 )
@@ -146,6 +148,7 @@ py_test(
         ":cluster",
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/keras:compat",
+        "//tensorflow_model_optimization/python/core/clustering/keras/experimental:cluster",
     ],
 )
 

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

@@ -81,8 +81,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)
@@ -93,8 +92,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
   }
 
   model = keras.Sequential([
@@ -128,6 +126,98 @@ def cluster_weights(to_cluster,
   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):
+  """Modify a keras layer or model to be clustered during training (private method).
+
+  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([
+      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: `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.
+
+  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.

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

@@ -26,6 +26,8 @@
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 from tensorflow_model_optimization.python.core.keras import compat
 
+from tensorflow_model_optimization.python.core.clustering.keras.experimental import cluster as experimental_cluster
+
 keras = tf.keras
 layers = keras.layers
 test = tf.test
@@ -159,7 +161,7 @@ def testSparsityIsPreservedDuringTraining(self):
         "preserve_sparsity": True
     }
 
-    clustered_model = cluster.cluster_weights(original_model, **clustering_params)
+    clustered_model = experimental_cluster.cluster_weights(original_model, **clustering_params)
 
     stripped_model_before_tuning = cluster.strip_clustering(clustered_model)
     weights_before_tuning = stripped_model_before_tuning.get_weights()[0]

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

@@ -26,6 +26,8 @@
 from tensorflow_model_optimization.python.core.clustering.keras import clusterable_layer
 from tensorflow_model_optimization.python.core.clustering.keras import clustering_registry
 
+from tensorflow_model_optimization.python.core.clustering.keras.experimental import cluster as experimental_cluster
+
 keras = tf.keras
 errors_impl = tf.errors
 layers = keras.layers
@@ -120,7 +122,7 @@ def testClusterKerasClusterableLayerWithSparsityPreservation(self):
     """
     preserve_sparsity_params = { 'preserve_sparsity': True }
     params = { **self.params, **preserve_sparsity_params }
-    wrapped_layer = cluster.cluster_weights(self.keras_clusterable_layer, **params)
+    wrapped_layer = experimental_cluster.cluster_weights(self.keras_clusterable_layer, **params)
 
     self._validate_clustered_layer(self.keras_clusterable_layer, wrapped_layer)
 
@@ -184,7 +186,7 @@ def testClusterCustomClusterableLayerWithSparsityPreservation(self):
     """
     preserve_sparsity_params = { 'preserve_sparsity': True }
     params = { **self.params, **preserve_sparsity_params }
-    wrapped_layer = cluster.cluster_weights(self.custom_clusterable_layer, **params)
+    wrapped_layer = experimental_cluster.cluster_weights(self.custom_clusterable_layer, **params)
     self.model.add(wrapped_layer)
     self.model.build(input_shape=(10, 1))
 
@@ -230,7 +232,7 @@ def testClusterSequentialModelSelectivelyWithSparsityPreservation(self):
     preserve_sparsity_params = { 'preserve_sparsity': True }
     params = { **self.params, **preserve_sparsity_params }
     clustered_model = keras.Sequential()
-    clustered_model.add(cluster.cluster_weights(self.keras_clusterable_layer, **params))
+    clustered_model.add(experimental_cluster.cluster_weights(self.keras_clusterable_layer, **params))
     clustered_model.add(self.keras_clusterable_layer)
     clustered_model.build(input_shape=(1, 10))
 
@@ -263,7 +265,7 @@ def testClusterFunctionalModelSelectivelyWithSparsityPreservation(self):
     params = { **self.params, **preserve_sparsity_params }
     i1 = keras.Input(shape=(10,))
     i2 = keras.Input(shape=(10,))
-    x1 = cluster.cluster_weights(layers.Dense(10), **params)(i1)
+    x1 = experimental_cluster.cluster_weights(layers.Dense(10), **params)(i1)
     x2 = layers.Dense(10)(i2)
     outputs = layers.Add()([x1, x2])
     clustered_model = keras.Model(inputs=[i1, i2], outputs=outputs)
@@ -302,7 +304,7 @@ def testClusterModelValidLayersSuccessfulWithSparsityPreservation(self):
         self.keras_non_clusterable_layer,
         self.custom_clusterable_layer
     ])
-    clustered_model = cluster.cluster_weights(model, **params)
+    clustered_model = experimental_cluster.cluster_weights(model, **params)
     clustered_model.build(input_shape=(1, 28, 28, 1))
 
     self.assertEqual(len(model.layers), len(clustered_model.layers))

tensorflow_model_optimization/python/core/clustering/keras/experimental/BUILD

@@ -0,0 +1,28 @@
+package(default_visibility = [
+    "//tensorflow_model_optimization:__subpackages__",
+])
+
+licenses(["notice"])  # Apache 2.0
+
+exports_files(["LICENSE"])
+
+py_library(
+    name = "experimental",
+    srcs = [
+        "__init__.py",
+    ],
+    srcs_version = "PY3",
+    deps = [
+        ":cluster",
+    ],
+)
+
+py_library(
+    name = "cluster",
+    srcs = ["cluster.py"],
+    srcs_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow_model_optimization/python/core/clustering/keras:cluster",
+    ],
+)

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

@@ -0,0 +1,109 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Experimental clustering API functions for Keras models."""
+
+from tensorflow_model_optimization.python.core.clustering.keras.cluster import _cluster_weights
+
+
+def cluster_weights(to_cluster,
+                    number_of_clusters,
+                    cluster_centroids_init,
+                    preserve_sparsity,
+                    **kwargs):
+  """Modify a keras layer or model to be clustered during training (experimental).
+
+  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:
+
+  ```python
+  clustering_params = {
+    'number_of_clusters': 8,
+    'cluster_centroids_init': CentroidInitialization.DENSITY_BASED,
+    'preserve_sparsity': True
+  }
+
+  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: `tfmot.clustering.keras.CentroidInitialization`
+        instance that determines how the cluster centroids will be initialized.
+      preserve_sparsity: 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.
+
+  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,
+                          **kwargs)