Merge pull request #951 from wwwind:clustering_per_channel_bugfix

PiperOrigin-RevId: 448168313

Author: tensorflower-gardener

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

@@ -154,7 +154,7 @@ def __init__(self,
     if hasattr(layer, 'data_format'):
       self.data_format = self.layer.data_format
     else:
-      self.data_format = None
+      self.data_format = 'channels_last'
 
     # Save the input shape specified in the build
     self.build_input_shape = None
@@ -215,7 +215,7 @@ def build(self, input_shape):
       centroid_init = centroid_init_factory.get_centroid_initializer(
           self.cluster_centroids_init)(weight, self.number_of_clusters,
                                        self.cluster_per_channel,
-                                       self.num_channels,
+                                       self.data_format,
                                        self.preserve_sparsity)
 
       # Init the cluster centroids

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

@@ -44,8 +44,15 @@ class AbstractCentroidsInitialisation:
   """
 
   def __init__(self, weights, number_of_clusters,
-               cluster_per_channel=False, data_format=None,
+               cluster_per_channel=False, data_format='channels_last',
                preserve_sparsity=False):
+
+    # Input checks
+    if (data_format != 'channels_first' and data_format != 'channels_last'):
+      raise ValueError(
+          'The given parameter data_format is not correct: {input}'.format(
+              input = data_format))
+
     self.weights = weights
     self.number_of_clusters = number_of_clusters
     self.cluster_per_channel = cluster_per_channel

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

@@ -66,6 +66,13 @@ def testThrowsValueErrorForNonExistingInit(self):
     with self.assertRaises(ValueError):
       self.factory.get_centroid_initializer("DEADBEEF")
 
+  def testThrowsValueErrorForWrongDataFormat(self):
+    """Verifies that the centroid initializer factory method raises an exception
+    when invoked with an wrong type of data_format."""
+    f = self.factory.get_centroid_initializer("CentroidInitialization.KMEANS_PLUS_PLUS")
+    with self.assertRaises(ValueError):
+      f([1, 2], 2, True, "NCHW")
+
   @parameterized.parameters(
       (0, 0, 1, 1, 1, 0),
       (0, 0, 5, 5, 1, 0),