Clusterable layer API.

Change-Id: I9bb372717923348d4c42fdecf7eeea64ce44f49b

Author: wwwind

tensorflow_model_optimization/python/core/clustering/keras/BUILD

@@ -164,3 +164,14 @@ py_test(
         "//tensorflow_model_optimization/python/core/keras:test_utils",
     ],
 )
+
+py_test(
+    name = "mnist_customerable_test",
+    srcs = ["mnist_customerable_test.py"],
+    python_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":cluster"
+        # tensorflow dep1,
+    ],
+)

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

@@ -55,6 +55,46 @@ class CustomNonClusterableLayer(layers.Dense):
   pass
 
 
+class MyCustomerableLayer(keras.layers.Dense,
+  clusterable_layer.ClusterableLayer):
+
+  def __init__(self, num_units):
+      super().__init__(num_units)
+
+  def get_clusterable_weights(self):
+    # Cluster kernel and bias.
+    return [('kernel', self.kernel), ('bias', self.bias)]
+
+class MyCustomerableLayerInvalid(keras.layers.Dense,
+  clusterable_layer.ClusterableLayer):
+  """ This layer is invalid, because it does not provide
+  get_clusterable_weights function.
+  """
+  def __init__(self, num_units):
+      super().__init__(num_units)
+
+class TestCustomerableWeightsCA(clustering_registry.AbstractClusteringAlgorithm):
+  """ Dummy class derived from AbstractClusteringAlgorithm."""
+  def get_pulling_indices(self, weight):
+    return [1, 2, 3]
+
+class KerasCustomLayerClusterable(keras.layers.Layer,
+  clusterable_layer.ClusterableLayer):
+  """ This keras custom layer is derived from ClusterableLayer
+  and it provides own implementation of the clustering
+  algorithm.
+  """
+
+  def __init__(self):
+    super().__init__()
+    self.kernel = None
+
+  def get_clusterable_weights(self):
+    return [('kernel', self.kernel)]
+
+  def get_clusterable_algorithm(self, weight_name):
+    return TestCustomerableWeightsCA
+
 class ClusterTest(test.TestCase, parameterized.TestCase):
   """Unit tests for the cluster module."""
 
@@ -67,6 +107,8 @@ def setUp(self):
     self.custom_clusterable_layer = CustomClusterableLayer(10)
     self.custom_non_clusterable_layer = CustomNonClusterableLayer(10)
     self.keras_depthwiseconv2d_layer = layers.DepthwiseConv2D((3, 3), (1, 1))
+    self.customerable_layer = MyCustomerableLayer(10)
+    self.keras_custom_layer = KerasCustomLayer()
 
     clustering_registry.ClusteringLookupRegistry.register_new_implementation(
         {
@@ -183,12 +225,54 @@ def testClusterCustomNonClusterableLayer(self):
       cluster_wrapper.ClusterWeights(custom_non_clusterable_layer,
                                      **self.params)
 
+  def testClusterMyCustomerableLayer(self):
+    # we have weights to cluster.
+    customerable_layer = self.customerable_layer
+    customerable_layer.build(input_shape=(10, 10))
+
+    wrapped_layer = cluster_wrapper.ClusterWeights(customerable_layer,
+                                     **self.params)
+
+    self.assertIsInstance(wrapped_layer, cluster_wrapper.ClusterWeights)
+
+  def testKerasCustomLayerClusterable(self):
+    """
+    Verifies that we can wrap keras custom layer that is customerable.
+    """
+    customerable_layer = KerasCustomLayerClusterable()
+    wrapped_layer = cluster_wrapper.ClusterWeights(customerable_layer,
+                                     **self.params)
+
+    self.assertIsInstance(wrapped_layer, cluster_wrapper.ClusterWeights)
+
+  def testClusterMyCustomerableLayerInvalid(self):
+    """
+    Verifies that assertion is thrown when function
+    get_clusterable_weights is not provided.
+    """
+    with self.assertRaises(TypeError):
+      MyCustomerableLayerInvalid(10) # pylint: disable=abstract-class-instantiated
+
+  def testClusterKerasCustomLayer(self):
+    """
+    Verifies that attempting to cluster a keras custom layer raises
+    an exception.
+    """
+    # If layer is not built, it has not weights, so
+    # we just skip it.
+    keras_custom_layer = self.keras_custom_layer
+    cluster_wrapper.ClusterWeights(keras_custom_layer,
+                                  **self.params)
+    # We need to build weights before check that clustering is not supported.
+    keras_custom_layer.build(input_shape=(10, 10))
+    with self.assertRaises(ValueError):
+      cluster_wrapper.ClusterWeights(keras_custom_layer,
+                                     **self.params)
+
+>>>>>>> 8fe29ec... MLTOOLS-1031 Customerable layer API.
   @keras_parameterized.run_all_keras_modes
   def testClusterSequentialModelSelectively(self):
-    """Verifies that layers within a sequential model can be clustered selectively."""
     clustered_model = keras.Sequential()
-    clustered_model.add(
-        cluster.cluster_weights(self.keras_clusterable_layer, **self.params))
     clustered_model.add(self.keras_clusterable_layer)
     clustered_model.build(input_shape=(1, 10))
 

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

@@ -178,7 +178,10 @@ def build(self, input_shape):
       # variable is either in the self._trainable_weights or in
       # self._non_trainable_weights and self.weights is the result of
       # concatenation of those arrays
-      original_index = self.layer.weights.index(weight)
+      original_index = 0
+      for i in range(len(self.layer.weights)):
+        if self.layer.weights[i].name == weight.name:
+          original_index = i
       self.gone_variables.append(original_index)
 
       # Again, not sure if this is needed. Leaving for now.

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

@@ -22,8 +22,13 @@
 class ClusterableLayer:
   """Abstract Base Class for making your own keras layer clusterable.
 
-  Custom keras layers that need to support clustering should implement this
-  class.
+  Your layer could be derived from a keras built-in layer or
+  it could be a keras custom layer.
+
+  The function get_clusterable_weights should be provided in both cases.
+
+  The function get_clusterable_algorithm is provided, when weights for
+  clustering is added in the keras layer.
 
   """
 
@@ -40,3 +45,21 @@ def get_clusterable_weights(self):
         kernel object itself.
     """
     raise NotImplementedError('Must be implemented in subclasses.')
+
+  def get_clusterable_algorithm(self, weight_name):
+    """Returns class with the clustering algorithm for the given weight_name.
+
+    This function needs to be implemented for the customerable layers.
+    If the layer is derived from the built-in keras layer, the clustering
+    algorithm for the base built-in keras layer is used.
+
+    The returned class should be derived from AbstractClusteringAlgorithm and
+    implements the function get_pulling_indices.
+    This function is used to provide a special lookup function for the custom weights.
+    It reshapes and tile centroids the same way as the weights. This allows us
+    to find pulling indices efficiently.
+
+    Args:
+        weight_name ([string]): The name of the weight variable.
+    """
+    return None

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

@@ -170,7 +170,6 @@ def get_pulling_indices(self, weight):
 
     return pulling_indices
 
-
 class ClusteringLookupRegistry(object):
   """
   The keys represent built-in keras layers and the values represent the
@@ -179,17 +178,17 @@ class ClusteringLookupRegistry(object):
   work on, or the strategy is not currently supported
   """
   _LAYERS_RESHAPE_MAP = {
-      layers.Conv1D: {'kernel': ConvolutionalWeightsCA},
-      layers.Conv2D: {'kernel': ConvolutionalWeightsCA},
-      layers.Conv2DTranspose: {'kernel': ConvolutionalWeightsCA},
-      layers.Conv3D: {'kernel': ConvolutionalWeightsCA},
-      layers.Conv3DTranspose: {'kernel': ConvolutionalWeightsCA},
-      layers.SeparableConv1D: {'pointwise_kernel': ConvolutionalWeightsCA},
-      layers.SeparableConv2D: {'pointwise_kernel': ConvolutionalWeightsCA},
-      layers.Dense: {'kernel': DenseWeightsCA},
-      layers.Embedding: {'embeddings': DenseWeightsCA},
-      layers.LocallyConnected1D: {'kernel': ConvolutionalWeightsCA},
-      layers.LocallyConnected2D: {'kernel': ConvolutionalWeightsCA},
+      layers.Conv1D: {'kernel': ConvolutionalWeightsCA, 'bias': BiasWeightsCA},
+      layers.Conv2D: {'kernel': ConvolutionalWeightsCA, 'bias': BiasWeightsCA},
+      layers.Conv2DTranspose: {'kernel': ConvolutionalWeightsCA, 'bias': BiasWeightsCA},
+      layers.Conv3D: {'kernel': ConvolutionalWeightsCA, 'bias': BiasWeightsCA},
+      layers.Conv3DTranspose: {'kernel': ConvolutionalWeightsCA, 'bias': BiasWeightsCA},
+      layers.SeparableConv1D: {'pointwise_kernel': ConvolutionalWeightsCA, 'bias': BiasWeightsCA},
+      layers.SeparableConv2D: {'pointwise_kernel': ConvolutionalWeightsCA, 'bias': BiasWeightsCA},
+      layers.Dense: {'kernel': DenseWeightsCA, 'bias': BiasWeightsCA},
+      layers.Embedding: {'embeddings': DenseWeightsCA, 'bias': BiasWeightsCA},
+      layers.LocallyConnected1D: {'kernel': ConvolutionalWeightsCA, 'bias': BiasWeightsCA},
+      layers.LocallyConnected2D: {'kernel': ConvolutionalWeightsCA, 'bias': BiasWeightsCA},
   }
 
   @classmethod
@@ -229,15 +228,37 @@ def get_clustering_impl(cls, layer, weight_name):
     :param weight_name: concrete weight name to be clustered.
     :return: a concrete implementation of a lookup algorithm
     """
+    custom_layer_of_built_layer = None
     if not layer.__class__ in cls._LAYERS_RESHAPE_MAP:
-      raise ValueError(
-          "Class {given_class} has not been registerd in the"
-          "ClusteringLookupRegistry. Use ClusteringLookupRegistry."
-          "register_new_implemenetation to fix this.".format(
-              given_class=layer.__class__
-          )
-      )
-    if weight_name not in cls._LAYERS_RESHAPE_MAP[layer.__class__]:
+      # Checks whether we have a customerable layer derived from built-in keras class.
+      for key in cls._LAYERS_RESHAPE_MAP:
+        if issubclass(layer.__class__, key):
+          custom_layer_of_built_layer = key
+      if not custom_layer_of_built_layer:
+        # Checks whether we have a customerable layer that provides
+        # clusterable algorithm for the given weights.
+        if issubclass(layer.__class__, clusterable_layer.ClusterableLayer) and \
+          layer.get_clusterable_algorithm is not None:
+            ans = layer.get_clusterable_algorithm(weight_name)
+            if not ans:
+              raise ValueError(
+                "Class {given_class} does not provided clustering algorithm"
+                "for the weights with the name {weight_name}.".format(
+                    given_class=layer.__class__, weight_name=weight_name
+                )
+              )
+            else:
+              return ans
+        else:
+          raise ValueError(
+            "Class {given_class} has not derived from ClusterableLayer"
+            "or the funtion get_pulling_indices is not provided.".format(
+                given_class=layer.__class__
+            )
+        )
+    else:
+      custom_layer_of_built_layer = layer.__class__
+    if weight_name not in cls._LAYERS_RESHAPE_MAP[custom_layer_of_built_layer]:
       raise ValueError(
           "Weight with the name '{given_weight_name}' for class {given_class} "
           "has not been registerd in the ClusteringLookupRegistry. Use "
@@ -249,7 +270,7 @@ def get_clustering_impl(cls, layer, weight_name):
       )
     # Different weights will have different shapes hence there is double hash
     # map lookup.
-    return cls._LAYERS_RESHAPE_MAP[layer.__class__][weight_name]
+    return cls._LAYERS_RESHAPE_MAP[custom_layer_of_built_layer][weight_name]
 
 
 class ClusteringRegistry(object):

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

@@ -201,9 +201,42 @@ def testGetClusteringImplFailsWithKnonwClassUnknownWeight(self):
       ClusteringLookupRegistry.get_clustering_impl(layers.Dense(10),
                                                    'no_such_weight')
 
+class KerasCustomLayerClusterableInvalid(keras.layers.Layer,
+  clusterable_layer.ClusterableLayer):
+  """ This keras custom layer is derived from ClusterableLayer
+  and it provides own implementation of the clustering
+  algorithm.
+  """
+
+  def __init__(self, units=10):
+    super(KerasCustomLayerClusterableInvalid, self).__init__()
+    self.units = units
+
+  def build(self, input_shape):
+    self.w = self.add_weight(
+      shape=(input_shape[-1], self.units),
+      initializer="random_normal",
+      trainable=True,
+    )
+
+  def get_clusterable_weights(self):
+    return [('w', self.w)]
+
+  def testKerasCustomLayerClusterableInvalid(self):
+    """
+    Verifies that get_clustering_impl() raises an error when invoked with a
+    keras custom layer derived from ClusterableLayer, but the function
+    get_clustering_algorithm is not provided.
+    """
+    with self.assertRaises(ValueError):
+      ClusteringLookupRegistry.get_clustering_impl(
+        KerasCustomLayerClusterableInvalid(), 'w')
+
   @parameterized.parameters(
       (layers.Conv2D, 'kernel', clustering_registry.ConvolutionalWeightsCA),
       (layers.Conv1D, 'kernel', clustering_registry.ConvolutionalWeightsCA),
+      (layers.Conv2D, 'bias', clustering_registry.BiasWeightsCA),
+      (layers.Conv1D, 'bias', clustering_registry.BiasWeightsCA),
   )
   def testReturnsResultsForKnownTypeKnownWeights(self,
                                                  layer_type,