[clustering] Clusterable layer API

[wwwind]

This PR extends our support for clustering the following two user cases:

1. The user wants to specify that the bias for the dense layer should be clustered as well.
   In this case the user needs to create a new layer derived from Dense and ClusterableLayer and
   provide what is needed to be clustered in the function get_clusterable_weights.

class MyDenseLayer(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)]

2. This user case is for a keras custom layer. We don't support a keras custom layer by default.
   To be able to cluster it, the user needs to derive it from ClusterableLayer and provides two functions:
   get_clusterable_weights, which specifies what should be clustered as in the first user case,
   get_clusterable_algorithm, which is needed to specify the layout of the weight tensor.
   The example is given below:

class MyClusterableLayer(keras.layers.Layer, clusterable_layer.ClusterableLayer):

  def __init__(self, units=32):
    super(MyClusterableLayer, 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,
    )
    self.b = self.add_weight(
      shape=(self.units,),
      initializer="random_normal",
      trainable=False
   )

  def call(self, inputs):
    return tf.matmul(inputs, self.w) + self.b

  def get_clusterable_weights(self):
    # Cluster only weights 'w'
    return [('w', self.w)]

  def get_clusterable_algorithm(self, weight_name):
     if weight_name == 'w':
       return ClusterableWeightsCA
     else:
       # We don't cluster other weights.
       return None

, where

class ClusterableWeightsCA(clustering_algorithm.AbstractClusteringAlgorithm):

  def get_pulling_indices(self, weight):
    clst_num = self.cluster_centroids.shape[0]
    tiled_weights = tf.tile(tf.expand_dims(weight, axis=2), [1, 1, clst_num])
    tiled_cluster_centroids = tf.tile(
        tf.reshape(self.cluster_centroids, [1, 1, clst_num]),
        [weight.shape[0], weight.shape[1], 1])

    # We find the nearest cluster centroids and store them so that ops can build
    # their kernels upon it
    pulling_indices = tf.argmin(tf.abs(tiled_weights - tiled_cluster_centroids),
                                axis=2)

    return pulling_indices

Both these user cases are added as a mnist test in the new test file: mnist_clusterable_layer_test.py.

[akarmi]

Thank you, and apologies for taking so long to review!

Can you please tidy up the description? The formatting is a bit off and makes it difficult to read. Also, the code snippet for the second use case does not include get_clusterable_algorithm().

[akarmi]

please, remove!

[akarmi]

nit: MyClusterableLayer instead? There are also similar instances in other part of the PR.

[akarmi]

We don't cluster biases by default. Will this change such a behaviour? It'd be good to have a test for this.

[akarmi]

nit: please rename to mnist_clusterable_layer_test.py

[akarmi]

Thanks. Looks good. Just a minor request.

[akarmi]

Please, add a similar assertion check for the number of unique bias values to the one for the weights above.

[akarmi]

Thanks

[akarmi]

Thank you!

[akarmi]

Apologies, I missed it initially. We need to expose ClusterableLayer in the public API, which will require exposing AbstractClusteringAlgorithm as well. The latter then needs moving out of the clustering_registry.py into its own module, to prevent exposing the implementation guts in the public API, and we should also then clean up clustering_registry.py to remove unused functions, e.g. register_new_implementation(). Any thoughts?

[wwwind]

Hi @akarmi Please re-review this PR - all comments are addressed. Thanks!

[akarmi]

Thanks.
@daverim, as this change affects the public API, please review it as well.

[akarmi]

Thanks

[akarmi]

@daverim, any feedback? - we would like to merge this now.

[wwwind]

Hi @daverim I updated this PR with the merged changes. I checked the test that is left in this PR locally and it works as expected.
Could we please merge this PR ? Thanks!