Merge pull request #616 from wwwind:clusterable_layer

PiperOrigin-RevId: 368158644

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/api/clustering/keras/__init__.py

@@ -19,7 +19,10 @@
 from tensorflow_model_optimization.python.core.clustering.keras.cluster import cluster_scope
 from tensorflow_model_optimization.python.core.clustering.keras.cluster import cluster_weights
 from tensorflow_model_optimization.python.core.clustering.keras.cluster import strip_clustering
+from tensorflow_model_optimization.python.core.clustering.keras import clusterable_layer
 
 from tensorflow_model_optimization.python.core.clustering.keras.cluster_config import CentroidInitialization
+from tensorflow_model_optimization.python.core.clustering.keras.clustering_algorithm import AbstractClusteringAlgorithm
 from tensorflow_model_optimization.python.core.clustering.keras.clustering_callbacks import ClusteringSummaries
+from tensorflow_model_optimization.python.core.clustering.keras.clusterable_layer import ClusterableLayer
 # pylint: enable=g-bad-import-order

tensorflow_model_optimization/python/core/clustering/keras/BUILD

@@ -44,6 +44,7 @@ py_library(
     visibility = ["//visibility:public"],
     deps = [
         ":clusterable_layer",
+        ":clustering_algorithm",
     ],
 )
 
@@ -74,6 +75,16 @@ py_library(
     ],
 )
 
+py_library(
+    name = "clustering_algorithm",
+    srcs = ["clustering_algorithm.py"],
+    srcs_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        # tensorflow dep1,
+    ],
+)
+
 py_library(
     name = "clustering_callbacks",
     srcs = ["clustering_callbacks.py"],
@@ -165,3 +176,25 @@ py_test(
         "//tensorflow_model_optimization/python/core/keras:test_utils",
     ],
 )
+
+py_test(
+    name = "mnist_clusterable_layer_test",
+    srcs = ["mnist_clusterable_layer_test.py"],
+    python_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":cluster",
+        # tensorflow dep1,
+    ],
+)
+
+py_test(
+    name = "mnist_clustering_test",
+    srcs = ["mnist_clustering_test.py"],
+    python_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":cluster",
+        # tensorflow dep1,
+    ],
+)

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

@@ -54,26 +54,64 @@ def get_clusterable_weights(self):
 class CustomNonClusterableLayer(layers.Dense):
   pass
 
+
 class KerasCustomLayer(keras.layers.Layer):
+
   def __init__(self, units=32):
     super(KerasCustomLayer, 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,
+        shape=(input_shape[-1], self.units),
+        initializer='random_normal',
+        trainable=True,
     )
     self.b = self.add_weight(
-      shape=(self.units,),
-      initializer="random_normal",
-      trainable=False
-    )
+        shape=(self.units,), initializer='random_normal', trainable=False)
 
   def call(self, inputs):
     return tf.matmul(inputs, self.w) + self.b
 
+
+class MyClusterableLayer(keras.layers.Dense,
+                         clusterable_layer.ClusterableLayer):
+
+  def get_clusterable_weights(self):
+    # Cluster kernel and bias.
+    return [('kernel', self.kernel), ('bias', self.bias)]
+
+
+class MyClusterableLayerInvalid(keras.layers.Dense,
+                                clusterable_layer.ClusterableLayer):
+  """This layer is invalid: it does not implement get_clusterable_weights(self).
+  """
+  pass
+
+
+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):
+  """Custom Keras clusterable layer, providing its own 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."""
 
@@ -87,6 +125,7 @@ def setUp(self):
     self.custom_non_clusterable_layer = CustomNonClusterableLayer(10)
     self.keras_depthwiseconv2d_layer = layers.DepthwiseConv2D((3, 3), (1, 1))
     self.keras_custom_layer = KerasCustomLayer()
+    self.clusterable_layer = MyClusterableLayer(10)
 
     clustering_registry.ClusteringLookupRegistry.register_new_implementation(
         {
@@ -204,24 +243,37 @@ def testClusterCustomNonClusterableLayer(self):
                                      **self.params)
 
   def testClusterKerasCustomLayer(self):
-    """
-    Verifies that attempting to cluster a keras custom layer raises
-    an exception.
-    """
+    """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)
+    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)
+      cluster_wrapper.ClusterWeights(keras_custom_layer, **self.params)
+
+  def testClusterMyClusterableLayer(self):
+    # we have weights to cluster.
+    layer = self.clusterable_layer
+    layer.build(input_shape=(10, 10))
+
+    wrapped_layer = cluster_wrapper.ClusterWeights(layer, **self.params)
+    self.assertIsInstance(wrapped_layer, cluster_wrapper.ClusterWeights)
+
+  def testKerasCustomLayerClusterable(self):
+    """Verifies that we can wrap keras custom layer that is customerable."""
+    layer = KerasCustomLayerClusterable()
+    wrapped_layer = cluster_wrapper.ClusterWeights(layer, **self.params)
+    self.assertIsInstance(wrapped_layer, cluster_wrapper.ClusterWeights)
+
+  def testClusterMyClusterableLayerInvalid(self):
+    """Verifies that an exception is raised when get_clusterable_weights() is not implemented."""
+    with self.assertRaises(TypeError):
+      MyClusterableLayerInvalid(10)  # pylint: disable=abstract-class-instantiated
 
   @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))

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,22 @@ def get_clusterable_weights(self):
         kernel object itself.
     """
     raise NotImplementedError('Must be implemented in subclasses.')
+
+  def get_clusterable_algorithm(self, weight_name):  # pylint: disable=unused-argument
+    """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_algorithm.py

@@ -0,0 +1,133 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Abstract base class for clustering algorithm."""
+
+import abc
+import six
+import tensorflow as tf
+
+
+@six.add_metaclass(abc.ABCMeta)
+class AbstractClusteringAlgorithm(object):
+  """Abstrac class to implement highly efficient vectorised look-ups.
+
+    We do not utilise looping for that purpose, instead we `smartly` reshape and
+    tile arrays. The trade-off is that we are potentially using way more memory
+    than we would have if looping is used.
+
+    Each class that inherits from this class is supposed to implement a
+    particular lookup function for a certain shape.
+
+    For example, look-ups for 2D table will be different in the case of 3D.
+  """
+
+  def __init__(self, clusters_centroids):
+    """Generating clustered tensors.
+
+    For generating clustered tensors we will need two things: cluster
+    centroids and the final shape tensor must have.
+
+    Args:
+      clusters_centroids: An array of shape (N,) that contains initial
+      values of clusters centroids.
+    """
+    self.cluster_centroids = clusters_centroids
+
+  @abc.abstractmethod
+  def get_pulling_indices(self, weight):
+    """Returns indices of closest cluster centroids.
+
+    Takes a weight(can be 1D, 2D or ND) and creates tf.int32 array of the
+    same shape that will hold indices of cluster centroids clustered arrays
+    elements will be pulled from.
+
+    In the current setup pulling indices are meant to be created once and
+    used everywhere
+
+    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
+    """
+    pass
+
+  @tf.custom_gradient
+  def add_custom_gradients(self, clst_weights, weights):
+    """Adds custom gradients in the backprop stage.
+
+    This function overrides gradients in the backprop stage: original mul
+    becomes add, tf.sign becomes tf.identity. It is to update the original
+    weights with the gradients updates directly from the layer wrapped. We
+    assume the gradients updates on individual elements inside a cluster
+    will be different so that there is no point of mapping the gradient
+    updates back to original weight matrix using the LUT.
+
+    Args:
+      clst_weights: cluster weights
+      weights: weights
+    Returns:
+      custom gradient
+    """
+    override_weights = tf.sign(tf.reshape(weights, shape=(-1,)) + 1e+6)
+    z = clst_weights * override_weights
+
+    def grad(dz):
+      return dz, dz
+
+    return z, grad
+
+  def get_clustered_weight(self, pulling_indices):
+    """Returns clustered weights.
+
+    Takes an array with integer number that represent lookup indices and
+    forms a new array according to the given indices.
+
+    Args:
+      pulling_indices: an array of indices used for lookup.
+    Returns:
+      An array with the same shape as `pulling_indices`. Each array element
+      is a member of self.cluster_centroids.
+    """
+
+    return tf.reshape(
+        tf.gather(self.cluster_centroids,
+                  tf.reshape(pulling_indices, shape=(-1,))),
+        shape=pulling_indices.shape)
+
+  def get_clustered_weight_forward(self, pulling_indices, weight):
+    """Returns clustered weights with custm gradiennt.
+
+    Takes indices (pulling_indices) and original weights (weight) as inputs
+
+    and then forms a new array according to the given indices. The original
+    weights (weight) here are added to the graph since we want the backprop
+    to update their values via the new implementation using tf.custom_gradient
+
+    Args:
+      pulling_indices: an array of indices used for lookup.
+      weight: the original weights of the wrapped layer.
+    Returns:
+      array with the same shape as `pulling_indices`. Each array element
+      is a member of self.cluster_centroids
+    """
+
+    x = tf.reshape(self.get_clustered_weight(pulling_indices), shape=(-1,))
+
+    return tf.reshape(
+        self.add_custom_gradients(x, tf.reshape(weight, shape=(-1,))),
+        pulling_indices.shape)