Support of clustering per channel in (P)CQAT.

Change-Id: I70d28b00f335b044e0fed3db7779a6cbbd56ded8

Author: wwwind

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

@@ -586,7 +586,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()
@@ -624,15 +624,27 @@ 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
+    assert layer_conv2d_kernel_weight is not None
+    nr_unique_weights = len(np.unique(layer_conv2d_kernel_weight.numpy()))
+    assert 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]))
+      assert 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 tf.transpose(
+          pulling_indices, perm=[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 tf.transpose(
+            pulling_indices, perm=[3, 0, 1, 2]))
+
     if self.cluster_gradient_aggregation == GradientAggregation.SUM:
       cluster_centroids = self.cluster_centroids
     elif self.cluster_gradient_aggregation == GradientAggregation.AVG:
@@ -268,7 +298,7 @@ def get_clustered_weight(self, pulling_indices, original_weight):
 
     for i in range(num_channels):
       clustered_weights.append(
-          tf.gather(cluster_centroids[i], pulling_indices[i]))
+        tf.gather(cluster_centroids[i], pulling_indices[i]))
 
     clustered_weight = tf.convert_to_tensor(clustered_weights)
 

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,43 +184,91 @@ 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]]]]))
+      ("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(
+      ("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)
+
+    # 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]]],
-        [[[0], [2]], [[1], [0]]]], [[1, 1, 0], [1, 1, 1]]),
+       [[[[0, 0, 0]], [[1, 1, 0]]],
+        [[[1, 0, 1]], [[0, 1, 0]]]],
+       [[1, 1], [1, 1], [1, 1]]),
       (GradientAggregation.SUM,
-       [[[[0], [0]], [[0], [1]]],
-        [[[0], [2]], [[1], [0]]]], [[3, 1, 0], [2, 1, 1]])
+       [[[[0, 0, 0]], [[1, 1, 0]]],
+        [[[1, 0, 1]], [[0, 1, 0]]]],
+       [[2, 2], [2, 2], [3, 1]])
       )
-  def testConvolutionalPerChannelCAGrad(self,
+  def testConvolutionalPerChannelCAGradChannelsLast(self,
                                         cluster_gradient_aggregation,
                                         pulling_indices,
                                         expected_grad_centroids):
-    """Verifies that the gradients of convolutional layer work as expected."""
+    """Verifies that the gradients of convolutional layer works."""
 
-    clustering_centroids = tf.Variable([[0., 1, 2], [3, 4, 5]],
+    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,
@@ -229,6 +277,37 @@ def testConvolutionalPerChannelCAGrad(self,
         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,
+        pulling_indices,
+        expected_grad_centroids,
+    )
 
 class CustomLayer(layers.Layer):
   """A custom non-clusterable layer class."""