Merge pull request #920 from wwwind:clustering_per_channel_cqat

PiperOrigin-RevId: 430602869

Author: tensorflower-gardener

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

@@ -562,7 +562,7 @@ def testMHA(self):
 
 class ClusterPerChannelIntegrationTest(tf.test.TestCase,
                                        parameterized.TestCase):
-  """Integration tests for per-channel clustering of Conv2D layer."""
+  """Integration test for cluster_per_channel for Conv2D layer."""
 
   def setUp(self):
     super(ClusterPerChannelIntegrationTest, self).setUp()
@@ -600,15 +600,28 @@ def testPerChannel(self):
         loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
         metrics=[tf.keras.metrics.SparseCategoricalAccuracy(name="accuracy")])
     clustered_model.fit(
-        self.x_train, self.y_train, epochs=1, batch_size=100, verbose=1)
+        self.x_train, self.y_train, epochs=2, batch_size=100, verbose=1)
 
     stripped_model = cluster.strip_clustering(clustered_model)
 
     layer_conv2d = stripped_model.layers[2]
+    layer_conv2d_kernel_weight = None
     for weight in layer_conv2d.weights:
       if "kernel" in weight.name:
-        nr_unique_weights = len(np.unique(weight.numpy()))
-        assert nr_unique_weights == self.nr_of_clusters*self.num_channels
+        layer_conv2d_kernel_weight = weight
+    self.assertIsNotNone(layer_conv2d_kernel_weight)
+    nr_unique_weights = len(np.unique(layer_conv2d_kernel_weight.numpy()))
+    self.assertEqual(nr_unique_weights, self.nr_of_clusters*self.num_channels)
+
+    # The above check is too general.
+    # We need to check that we actually have nr_of_clusters per channel.
+    # If more general case passed, then we do tighter check.
+    # Note that we assume that data_format is NHWC.
+    for i in range(self.num_channels):
+      nr_unique_weights_per_channel = len(
+          np.unique(layer_conv2d_kernel_weight[:, :, :, i]))
+      self.assertEqual(nr_unique_weights_per_channel, self.nr_of_clusters)
+
 
 if __name__ == "__main__":
   test.main()

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

@@ -53,7 +53,7 @@ def __init__(
         clusters centroids.
       cluster_gradient_aggregation: An enum that specify the aggregation method
         of the cluster gradient.
-      data_format: To be used in Per-Channel clustering to ensure the weight
+      data_format: To be used in cluster_per_channel to ensure the weight
         kernel is permuted properly when updating the weights and calculating
         gradients
     """
@@ -194,10 +194,24 @@ def get_clustered_weight(self, pulling_indices, original_weight):
     return clustered_weight
 
 
-class PerChannelCA(ClusteringAlgorithm):
+class ClusteringAlgorithmPerChannel(ClusteringAlgorithm):
   """Class for Per-channel clustering of Conv2D layers."""
 
   def get_pulling_indices(self, weight):
+    """Returns indices of closest cluster centroids.
+
+    This function is based on the function get_pulling_indices
+    of the base class ClusteringAlgorithm. We apply each per
+    channel of the convolutional layer.
+
+    Args:
+      weight: ND array of weights. For each weight in this array the closest
+        cluster centroids is found.
+
+    Returns:
+      ND array of the same shape as `weight` parameter of the type
+      tf.int32. The returned array contain weight lookup indices.
+    """
     channel_indices = []
 
     num_channels = (weight.shape[1] if self.data_format == "channels_first"
@@ -213,7 +227,13 @@ def get_pulling_indices(self, weight):
 
       channel_indices.append(pulling_indices)
 
-    return tf.convert_to_tensor(channel_indices)
+    pulling_indices = tf.convert_to_tensor(channel_indices)
+    pulling_indices = tf.transpose(
+        pulling_indices,
+        perm=(1, 0, 2, 3) if self.data_format == "channels_first" else
+        (1, 2, 3, 0))
+
+    return pulling_indices
 
   def get_clustered_weight(self, pulling_indices, original_weight):
     """Returns clustered weights with custom gradients.
@@ -240,6 +260,16 @@ def get_clustered_weight(self, pulling_indices, original_weight):
         original_weight.shape[1]
         if self.data_format == "channels_first" else original_weight.shape[-1])
 
+    # In case of channels_last, we have NHWC.
+    # In case of channels_first, we have NCHW.
+    # We need to transpose the tensor, so C is the first dimension
+    # and then we could loop over channels
+    pulling_indices = (
+        tf.transpose(
+            pulling_indices,
+            perm=(1, 0, 2, 3) if self.data_format == "channels_first" else
+            (3, 0, 1, 2)))
+
     if self.cluster_gradient_aggregation == GradientAggregation.SUM:
       cluster_centroids = self.cluster_centroids
     elif self.cluster_gradient_aggregation == GradientAggregation.AVG:

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

@@ -21,7 +21,7 @@
 
 layers = tf.keras.layers
 ClusteringAlgorithm = clustering_algorithm.ClusteringAlgorithm
-PerChannelCA = clustering_algorithm.PerChannelCA
+ClusteringAlgorithmPerChannel = clustering_algorithm.ClusteringAlgorithmPerChannel
 
 
 class ClusteringLookupRegistry(object):
@@ -43,7 +43,7 @@ def get_clustering_impl(cls, layer, weight_name, cluster_per_channel=False):
     # Per-channel clustering is only applied if the layer is a Conv2D,
     # ignored otherwise
     if cluster_per_channel and isinstance(layer, tf.keras.layers.Conv2D):
-      return PerChannelCA
+      return ClusteringAlgorithmPerChannel
 
     # Clusterable layer could provide own implementation of get_pulling_indices
     if (issubclass(layer.__class__, clusterable_layer.ClusterableLayer) and

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

@@ -184,44 +184,111 @@ def testConvolutionalWeightsCA(self, clustering_centroids, pulling_indices,
     self._check_pull_values(clustering_algo, pulling_indices, expected_output)
 
   @parameterized.parameters(
-      ([[0., 1, 2], [3, 4, 5]],
-       [[[[0], [0]], [[0], [1]]],
-        [[[0], [2]], [[1], [0]]]],
-       [[[[0], [0]], [[0], [0]]],
-        [[[0], [0]], [[1], [1]]]]))
-  def testConvolutionalWeightsPerChannelCA(self,
+      (
+          'channels_last',
+          [[1, 2], [3, 4], [5, 6]],  # 3 channels and 2 cluster per channel
+          # pulling indices has shape (2, 2, 1, 3)
+          [[[[0, 1, 0]], [[0, 1, 1]]], [[[1, 0, 1]], [[0, 1, 0]]]],
+          [[[[1, 4, 5]], [[1, 4, 6]]], [[[2, 3, 6]], [[1, 4, 5]]]]),
+      (
+          'channels_first',
+          [[1, 2], [3, 4], [4, 5], [6, 7]
+          ],  # 4 channels and 2 clusters per channel
+          # pulling indices has shape (1, 4, 2, 2)
+          [[[[0, 1], [1, 1]], [[0, 0], [0, 1]], [[1, 0], [0, 0]],
+            [[1, 1], [0, 0]]]],
+          [[[[1, 2], [2, 2]], [[3, 3], [3, 4]], [[5, 4], [4, 4]],
+            [[7, 7], [6, 6]]]]))
+  def testConvolutionalWeightsPerChannelCA(self, data_format,
                                            clustering_centroids,
                                            pulling_indices,
                                            expected_output):
-    """Verifies that PerChannelCA works as expected."""
+    """Verifies that get_clustered_weight function works as expected."""
     clustering_centroids = tf.Variable(clustering_centroids, dtype=tf.float32)
-    clustering_algo = clustering_registry.PerChannelCA(
-        clustering_centroids, GradientAggregation.SUM
+    clustering_algo = clustering_registry.ClusteringAlgorithmPerChannel(
+        clustering_centroids, GradientAggregation.SUM, data_format
     )
+    # Note that clustered_weights has the same shape as pulling_indices,
+    # because they are defined inside of the check function.
     self._check_pull_values(clustering_algo, pulling_indices, expected_output)
 
   @parameterized.parameters(
-      (GradientAggregation.AVG,
-       [[[[0], [0]], [[0], [1]]],
-        [[[0], [2]], [[1], [0]]]], [[1, 1, 0], [1, 1, 1]]),
-      (GradientAggregation.SUM,
-       [[[[0], [0]], [[0], [1]]],
-        [[[0], [2]], [[1], [0]]]], [[3, 1, 0], [2, 1, 1]])
-      )
-  def testConvolutionalPerChannelCAGrad(self,
-                                        cluster_gradient_aggregation,
-                                        pulling_indices,
-                                        expected_grad_centroids):
-    """Verifies that the gradients of convolutional layer work as expected."""
+      (
+          'channels_last',
+          [[1, 2], [3, 4], [5, 6]],  # 3 channels and 2 cluster per channel
+          # weight has shape (2, 2, 1, 3)
+          [[[[1.1, 3.2, 5.2]], [[2.0, 4.1, 5.2]]],
+           [[[2.1, 2., 6.1]], [[1., 5., 5.]]]],
+          # expected pulling indices
+          [[[[0, 0, 0]], [[1, 1, 0]]], [[[1, 0, 1]], [[0, 1, 0]]]]),
+      (
+          'channels_first',
+          # 4 channels and 2 clusters per channel
+          [[1, 2], [3, 4], [4, 5], [6, 7]],
+          # weight has shape (1, 4, 2, 2)
+          [[[[0.1, 1.5], [2.0, 1.1]], [[0., 3.5], [4.4, 4.]],
+            [[4.1, 4.2], [5.3, 6.]], [[7., 7.1], [6.1, 5.8]]]],
+          # expected pulling indices
+          [[[[0, 0], [1, 0]], [[0, 0], [1, 1]], [[0, 0], [1, 1]],
+            [[1, 1], [0, 0]]]]))
+  def testConvolutionalPullingIndicesPerChannelCA(self, data_format,
+                                                  clustering_centroids, weight,
+                                                  expected_output):
+    """Verifies that get_pulling_indices function works as expected."""
+    clustering_centroids = tf.Variable(clustering_centroids, dtype=tf.float32)
+    clustering_algo = clustering_registry.ClusteringAlgorithmPerChannel(
+        clustering_centroids, GradientAggregation.SUM, data_format
+    )
+    weight = tf.convert_to_tensor(weight)
+    pulling_indices = clustering_algo.get_pulling_indices(weight)
 
-    clustering_centroids = tf.Variable([[0., 1, 2], [3, 4, 5]],
+    # check that pulling_indices has the same shape as weight
+    self.assertEqual(pulling_indices.shape, weight.shape)
+    self.assertAllEqual(pulling_indices, expected_output)
+
+  @parameterized.parameters(
+      (GradientAggregation.AVG, [
+          [[[0, 0, 0]], [[1, 1, 0]]], [[[1, 0, 1]], [[0, 1, 0]]]],
+       [[1, 1], [1, 1], [1, 1]]),
+      (GradientAggregation.SUM, [
+          [[[0, 0, 0]], [[1, 1, 0]]], [[[1, 0, 1]], [[0, 1, 0]]]],
+       [[2, 2], [2, 2], [3, 1]]))
+  def testConvolutionalPerChannelCAGradChannelsLast(
+      self, cluster_gradient_aggregation, pulling_indices,
+      expected_grad_centroids):
+    """Verifies that the gradients of convolutional layer works."""
+
+    clustering_centroids = tf.Variable([[1, 2], [3, 4], [5, 6]],
                                        dtype=tf.float32)
-    weight = tf.constant([[[[0.1, 3.0]], [[0.2, 0.1]]],
-                          [[[0.1, 3.0]], [[0.2, 0.1]]]])
+    weight = tf.constant([[[[1.1, 3.2, 5.2]], [[2.0, 4.1, 5.2]]],
+                          [[[2.1, 2., 6.1]], [[1., 5., 5.]]]])
 
-    clustering_algo = clustering_registry.PerChannelCA(
-        clustering_centroids, cluster_gradient_aggregation
+    clustering_algo = clustering_registry.ClusteringAlgorithmPerChannel(
+        clustering_centroids, cluster_gradient_aggregation, 'channels_last')
+    self._check_gradients_clustered_weight(
+        clustering_algo,
+        weight,
+        pulling_indices,
+        expected_grad_centroids,
     )
+
+  @parameterized.parameters((GradientAggregation.AVG, [
+      [[[0, 0], [1, 0]], [[0, 0], [1, 1]], [[0, 0], [1, 1]], [[1, 1], [0, 0]]]
+  ], [[1, 1], [1, 1], [1, 1], [1, 1]]), (GradientAggregation.SUM, [
+      [[[0, 0], [1, 0]], [[0, 0], [1, 1]], [[0, 0], [1, 1]], [[1, 1], [0, 0]]]
+  ], [[3, 1], [2, 2], [2, 2], [2, 2]]))
+  def testConvolutionalPerChannelCAGradChannelsFirst(
+      self, cluster_gradient_aggregation, pulling_indices,
+      expected_grad_centroids):
+    """Verifies that the gradients of convolutional layer works."""
+
+    clustering_centroids = tf.Variable([[1, 2], [3, 4], [4, 5], [6, 7]],
+                                       dtype=tf.float32)
+    weight = tf.constant([[[[0.1, 1.5], [2.0, 1.1]], [[0., 3.5], [4.4, 4.]],
+                           [[4.1, 4.2], [5.3, 6.]], [[7., 7.1], [6.1, 5.8]]]])
+
+    clustering_algo = clustering_registry.ClusteringAlgorithmPerChannel(
+        clustering_centroids, cluster_gradient_aggregation, 'channels_first')
     self._check_gradients_clustered_weight(
         clustering_algo,
         weight,

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve/BUILD

@@ -37,6 +37,7 @@ py_strict_library(
     deps = [
         ":cluster_utils",
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config",
         "//tensorflow_model_optimization/python/core/clustering/keras:clustering_registry",
         "//tensorflow_model_optimization/python/core/quantization/keras:quant_ops",
         "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",