Merge pull request #986 from jamwar01:patch_issue_979

PiperOrigin-RevId: 460333065

Author: tensorflower-gardener

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

@@ -14,6 +14,8 @@
 # ==============================================================================
 """Clustering API functions for Keras models."""
 
+import warnings
+
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
@@ -266,6 +268,14 @@ def _add_clustering_wrapper(layer):
           **kwargs,
       )
 
+    # Skip clustering if Conv2D layer has insufficient number of weights
+    # for type of clustering
+    if isinstance(
+        layer,
+        tf.keras.layers.Conv2D) and not layer_has_enough_weights_to_cluster(
+            layer, number_of_clusters, cluster_per_channel):
+      return layer
+
     return cluster_wrapper.ClusterWeights(layer, number_of_clusters,
                                           cluster_centroids_init,
                                           preserve_sparsity,
@@ -355,3 +365,43 @@ def _strip_clustering_wrapper(layer):
   # Just copy the model with the right callback
   return tf.keras.models.clone_model(
       model, input_tensors=None, clone_function=_strip_clustering_wrapper)
+
+
+def layer_has_enough_weights_to_cluster(layer, number_of_clusters,
+                                        cluster_per_channel):
+  """Returns whether layer has enough weights to cluster.
+
+  Returns True if Conv2D layer has sufficient number of
+  weights to implement clustering, given an input number of clusters.
+
+  Args:
+    layer: input layer to return quantize configs for.
+    number_of_clusters: A number of cluster centroids to form clusters.
+    cluster_per_channel: An optional boolean value.
+  """
+  if not isinstance(layer, tf.keras.layers.Conv2D):
+    raise ValueError(f'Input layer should be Conv2D layer: {layer.name} given.')
+
+  if not layer.trainable_weights:
+    raise ValueError(f'Layer {layer.name} has no weights to cluster.')
+
+  number_of_layer_weights = tf.cast(tf.size(getattr(layer, 'kernel')), tf.int32)
+  channel_idx = 1 if layer.data_format == 'channels_first' else -1
+  number_of_channels = tf.size(layer.trainable_weights[channel_idx])
+
+  if cluster_per_channel:
+    weights_to_cluster = number_of_layer_weights / number_of_channels
+  else:
+    weights_to_cluster = number_of_layer_weights
+
+  if weights_to_cluster <= number_of_clusters:
+    has_enough_weights = False
+  else:
+    has_enough_weights = True
+
+  if not has_enough_weights:
+    warnings.warn(
+        f"Layer {layer.name} does not have enough weights to implement"
+        f"{'per-channel ' if cluster_per_channel else ''}clustering."
+        f" \nNo clustering was implemented for this layer.\n")
+  return has_enough_weights

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

@@ -288,6 +288,74 @@ def do_checks(layer, layer_name):
     do_checks(clustered_model.layers[2], "conv1d")
     do_checks(clustered_model.layers[3], "conv1d_transpose")
 
+  @parameterized.parameters(
+      (False, 16),  # number_of_clusters > Conv2D filters
+      (True,
+       8),  # number_of_clusters < Conv2D filters (but clustering by channel)
+      (True, 12),  # number_of_clusters = Conv2D filters
+      (False, 12),  # number_of_clusters = Conv2D filters
+  )
+  def testEndToEnd1x1Conv2d(self, cluster_per_channel, number_of_clusters):
+    """Test End to End clustering - model with 1x1 Conv2D.
+
+    Clustering should not be performed at all, since number of
+    weights in the layer is too low in all of these cases.
+
+    Args:
+      cluster_per_channel: An optional boolean value.
+      number_of_clusters: A number of cluster centroids to form clusters.
+    """
+    kernel_size = (1, 1)
+
+    inp = keras.layers.Input(shape=(28, 28), batch_size=16)
+    x = keras.layers.Reshape(target_shape=(28, 28, 1))(inp)
+    x = keras.layers.Conv2D(
+        filters=12, kernel_size=kernel_size, activation=tf.nn.relu)(
+            x)
+    model = keras.models.Model(inputs=inp, outputs=[x])
+
+    cluster_params = {
+        "number_of_clusters": number_of_clusters,
+        "cluster_per_channel": cluster_per_channel
+    }
+
+    # Get unique kernel weights on original model for comparison
+    original_unique_weights = model.layers[2].weights[0]
+
+    def apply_clustering(layer):
+      if isinstance(layer, keras.layers.Conv2D):
+        return cluster.cluster_weights(layer, **cluster_params)
+      return layer
+
+    # Ensure a warning is given to the user that clustering is not
+    # implemented for this layer
+    with self.assertWarnsRegex(Warning,
+                               r"Layer conv2d does not have enough weights"):
+      model_to_cluster = keras.models.clone_model(
+          model,
+          clone_function=apply_clustering,
+      )
+
+    model_to_cluster.compile(
+        loss=keras.losses.categorical_crossentropy,
+        optimizer="adam",
+        metrics=["accuracy"])
+    model_to_cluster.fit(
+        np.random.randn(*self._batch(model.input.get_shape().as_list(), 16)),
+        np.random.randn(*self._batch(model.output.get_shape().as_list(), 16)),
+        steps_per_epoch=1)
+    clustered_model = cluster.strip_clustering(model_to_cluster)
+
+    def do_checks(layer, layer_name, original_unique_weights):
+      self.assertEqual(layer.name, layer_name)
+      unique_weights = layer.weights[0]
+
+      # Ensure clustering was not performed on the 1x1 Conv
+      # (weights are identical to original unclustered layer)
+      self.assertAllEqual(unique_weights, original_unique_weights)
+
+    do_checks(clustered_model.layers[2], "conv2d", original_unique_weights)
+
   def testStripClusteringSequentialModelWithRegulariser(self):
     """Verifies that stripping the clustering wrappers from a sequential model produces the expected config."""
     original_model = keras.Sequential([

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

@@ -145,10 +145,10 @@ def setUp(self):
     }
 
   def _build_clustered_layer_model(self, layer, input_shape=(10, 1)):
-    wrapped_layer = cluster.cluster_weights(layer, **self.params)
-    self.model.add(wrapped_layer)
-    self.model.build(input_shape=input_shape)
-
+    self.model.add(keras.Input(shape=input_shape))
+    self.model.add(layer)
+    self.model.build()
+    wrapped_layer = cluster.cluster_weights(self.model.layers[0], **self.params)
     return wrapped_layer
 
   def _validate_clustered_layer(self, original_layer, wrapped_layer):
@@ -194,7 +194,7 @@ def testClusterKerasNonClusterableLayer(self):
   def testDepthwiseConv2DLayerNonClusterable(self):
     """Verifies that we don't cluster a DepthwiseConv2D layer, because clustering of this type of layer gives big unrecoverable accuracy loss."""
     wrapped_layer = self._build_clustered_layer_model(
-        self.keras_depthwiseconv2d_layer, input_shape=(1, 10, 10, 10))
+        self.keras_depthwiseconv2d_layer, input_shape=(10, 10, 10))
 
     self._validate_clustered_layer(self.keras_depthwiseconv2d_layer,
                                    wrapped_layer)
@@ -203,7 +203,7 @@ def testDepthwiseConv2DLayerNonClusterable(self):
   @keras_parameterized.run_all_keras_modes
   def testDenseLayer(self):
     """Verifies that we can cluster a Dense layer."""
-    input_shape = (4, 28, 1)
+    input_shape = (28, 1)
     wrapped_layer = self._build_clustered_layer_model(
         self.keras_dense_layer,
         input_shape=input_shape
@@ -217,7 +217,7 @@ def testDenseLayer(self):
   @keras_parameterized.run_all_keras_modes
   def testConv1DLayer(self):
     """Verifies that we can cluster a Conv1D layer."""
-    input_shape = (4, 28, 1)
+    input_shape = (28, 1)
     wrapped_layer = self._build_clustered_layer_model(
         self.keras_conv1d_layer,
         input_shape=input_shape)
@@ -230,7 +230,7 @@ def testConv1DLayer(self):
   @keras_parameterized.run_all_keras_modes
   def testConv1DTransposeLayer(self):
     """Verifies that we can cluster a Conv1DTranspose layer."""
-    input_shape = (4, 28, 1)
+    input_shape = (28, 1)
     wrapped_layer = self._build_clustered_layer_model(
         self.keras_conv1d_tr_layer,
         input_shape=input_shape)
@@ -243,7 +243,7 @@ def testConv1DTransposeLayer(self):
   @keras_parameterized.run_all_keras_modes
   def testConv2DLayer(self):
     """Verifies that we can cluster a Conv2D layer."""
-    input_shape = (4, 28, 28, 1)
+    input_shape = (28, 28, 1)
     wrapped_layer = self._build_clustered_layer_model(
         self.keras_conv2d_layer,
         input_shape=input_shape)
@@ -256,7 +256,7 @@ def testConv2DLayer(self):
   @keras_parameterized.run_all_keras_modes
   def testConv2DTransposeLayer(self):
     """Verifies that we can cluster a Conv2DTranspose layer."""
-    input_shape = (4, 28, 28, 1)
+    input_shape = (28, 28, 1)
     wrapped_layer = self._build_clustered_layer_model(
         self.keras_conv2d_tr_layer,
         input_shape=input_shape)
@@ -269,7 +269,7 @@ def testConv2DTransposeLayer(self):
   @keras_parameterized.run_all_keras_modes
   def testConv3DLayer(self):
     """Verifies that we can cluster a Conv3D layer."""
-    input_shape = (4, 28, 28, 28, 1)
+    input_shape = (28, 28, 28, 1)
     wrapped_layer = self._build_clustered_layer_model(
         self.keras_conv3d_layer,
         input_shape=input_shape)
@@ -732,7 +732,7 @@ def testClusterWeightsStrippedWeights(self):
   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."""
     i1 = keras.Input(shape=(1, 1, 1))
-    x1 = layers.Conv2D(1, 1)(i1)
+    x1 = layers.Conv2D(12, 1)(i1)
     outputs = x1
     model = keras.Model(inputs=[i1], outputs=outputs)
 

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve/cluster_preserve_integration_test.py

@@ -103,15 +103,20 @@ def _get_clustered_model(self, preserve_sparsity):
 
     return clustered_model
 
-  def _get_conv_model(self, nr_of_channels, data_format=None):
+  def _get_conv_model(self,
+                      nr_of_channels,
+                      data_format=None,
+                      kernel_size=(3, 3)):
     """Returns functional model with Conv2D layer."""
     inp = tf.keras.layers.Input(shape=(32, 32), batch_size=100)
     shape = (1, 32, 32) if data_format == 'channels_first' else (32, 32, 1)
     x = tf.keras.layers.Reshape(shape)(inp)
     x = tf.keras.layers.Conv2D(
-        filters=nr_of_channels, kernel_size=(3, 3),
+        filters=nr_of_channels,
+        kernel_size=kernel_size,
         data_format=data_format,
-        activation='relu')(x)
+        activation='relu')(
+            x)
     x = tf.keras.layers.MaxPool2D(2, 2)(x)
     out = tf.keras.layers.Flatten()(x)
     model = tf.keras.Model(inputs=inp, outputs=out)
@@ -130,11 +135,15 @@ def _compile_and_fit_conv_model(self, model, nr_epochs=1):
 
     return model
 
-  def _get_conv_clustered_model(self, nr_of_channels, nr_of_clusters,
-                                data_format, preserve_sparsity):
+  def _get_conv_clustered_model(self,
+                                nr_of_channels,
+                                nr_of_clusters,
+                                data_format,
+                                preserve_sparsity,
+                                kernel_size=(3, 3)):
     """Returns clustered per channel model with Conv2D layer."""
     tf.random.set_seed(42)
-    model = self._get_conv_model(nr_of_channels, data_format)
+    model = self._get_conv_model(nr_of_channels, data_format, kernel_size)
 
     if preserve_sparsity:
       # Make the convolutional layer sparse by nullifying half of weights
@@ -475,6 +484,81 @@ def testEndToEndPCQATClusteredPerChannel(self, data_format='channels_last'):
     cqat_sparsity = self._get_sparsity(stripped_cqat_model)
     self.assertLessEqual(cqat_sparsity[0], control_sparsity[0])
 
+  def testEndToEndPCQATClusteredPerChannelConv2d1x1(self,
+                                                    data_format='channels_last'
+                                                    ):
+    """Runs PCQAT for model containing a 1x1 Conv2D.
+
+    (with insufficient number of weights per channel).
+
+    Args:
+      data_format: Format of input data.
+    """
+    nr_of_channels = 12
+    nr_of_clusters = 4
+
+    # Ensure a warning is given to the user that
+    # clustering is not implemented for this layer
+    with self.assertWarnsRegex(Warning,
+                               r'Layer conv2d does not have enough weights'):
+      clustered_model = self._get_conv_clustered_model(
+          nr_of_channels,
+          nr_of_clusters,
+          data_format,
+          preserve_sparsity=True,
+          kernel_size=(1, 1))
+      stripped_model = cluster.strip_clustering(clustered_model)
+
+    # Save the kernel weights
+    conv2d_layer = stripped_model.layers[2]
+    self.assertEqual(conv2d_layer.name, 'conv2d')
+
+    for weight in conv2d_layer.weights:
+      if 'kernel' in weight.name:
+        # Original number of unique weights
+        nr_original_weights = len(np.unique(weight.numpy()))
+        self.assertLess(nr_original_weights, nr_of_channels * nr_of_clusters)
+
+        # Demonstrate unmodified test layer has less weights
+        # than requested clusters
+        for channel in range(nr_of_channels):
+          channel_weights = (
+              weight[:, channel, :, :]
+              if data_format == 'channels_first' else weight[:, :, :, channel])
+          nr_channel_weights = len(channel_weights)
+          self.assertGreater(nr_channel_weights, 0)
+          self.assertLessEqual(nr_channel_weights, nr_of_clusters)
+
+    # get sparsity before PCQAT training
+    # we expect that only one value will be returned
+    control_sparsity = self._get_sparsity(stripped_model)
+    self.assertGreater(control_sparsity[0], 0.5)
+
+    quant_aware_annotate_model = (
+        quantize.quantize_annotate_model(stripped_model))
+
+    with self.assertWarnsRegex(
+        Warning, r'No clustering performed on layer quant_conv2d'):
+      quant_aware_model = quantize.quantize_apply(
+          quant_aware_annotate_model,
+          scheme=default_8bit_cluster_preserve_quantize_scheme
+          .Default8BitClusterPreserveQuantizeScheme(preserve_sparsity=True))
+
+    # Lets train for more epochs to have a chance to scatter clusters
+    model = self._compile_and_fit_conv_model(quant_aware_model, 3)
+
+    stripped_cqat_model = strip_clustering_cqat(model)
+
+    # Check the unique weights of a certain layer of
+    # clustered_model and cqat_model, ensuring unchanged
+    layer_nr = 3
+    num_of_unique_weights_cqat = self._get_number_of_unique_weights(
+        stripped_cqat_model, layer_nr, 'kernel')
+    self.assertEqual(num_of_unique_weights_cqat, nr_original_weights)
+
+    cqat_sparsity = self._get_sparsity(stripped_cqat_model)
+    self.assertLessEqual(cqat_sparsity[0], control_sparsity[0])
+
   def testPassingNonPrunedModelToPCQAT(self):
     """Runs PCQAT as CQAT if the input model is not pruned."""
     preserve_sparsity = False

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve/cluster_preserve_quantize_registry.py

@@ -15,6 +15,7 @@
 """Registry responsible for built-in keras classes."""
 
 import logging
+import warnings
 
 import tensorflow as tf
 
@@ -259,6 +260,7 @@ def apply_cluster_preserve_quantize_config(self, layer, quantize_config):
     if self._no_trainable_weights(layer) or self._disable_cluster_preserve(
         layer):
       return quantize_config
+
     # Example: Conv2D, Dense layers
     if quantize_config.__class__.__name__ in self._LAYERS_CONFIG_MAP[
         layer.__class__].quantize_config_attrs:
@@ -277,11 +279,6 @@ class Default8bitClusterPreserveQuantizeRegistry(
     ClusterPreserveQuantizeRegistry):
   """Default 8 bit ClusterPreserveQuantizeRegistry."""
 
-  def __init__(self, preserve_sparsity):
-    super(Default8bitClusterPreserveQuantizeRegistry, self).__init__(
-        preserve_sparsity)
-    self.preserve_sparsity = preserve_sparsity
-
   def get_quantize_config(self, layer):
     """Returns the quantization config with weight_quantizer for a given layer.
 
@@ -364,6 +361,13 @@ def _build_clusters(self, name, layer):
     # Prepare clustering variables for the Keras graph when clusters
     # exist, assuming we do not use number_of_clusters larger than 1024
     if num_centroids > 1024:
+      warnings.warn(f'No clustering performed on layer {layer.name}.\n'
+                    f'Too many centroids to cluster.')
+      return result
+    # If not enough clusters, we do not preserve clustering
+    elif num_centroids <= 1:
+      warnings.warn(f'No clustering performed on layer {layer.name}.\n'
+                    f'Perhaps too many clusters requested for this layer?')
       return result
     else:
       clst_centroids_tf = layer.add_weight(