Model clustering for Model Optimization Toolkit implementation

[akarmi]

Motivation

This PR introduces a model clustering technique for Model Optimization
Toolkit. In this context the word "clustering" means when model
weights have only a few unique values. This allows for certain types
of hardware to benefit from advanced weight compression techniques
and the associated reduction in model memory footprint and bandwidth.
Clustering can also be combined with pruning and quantization reducing
network’s storage requirements even further.

Approach description

The technique description that is implemented can be found here:
https://arxiv.org/abs/1510.00149

From practical standpoint this is implemented using a lookup table to
hold the cluster centroid values during model training. The weight
array is populated with 'gather' operation so that during back
propagation the gradients can be calculated in a normal way. The lookup
table is then adjusted using the cumulative gradient values for the
weights that correspond to the same centroid.

The API is intended to be used in the same way pruning API is used.
The interfaces is intentionally made very similar to the one described
here:
https://www.tensorflow.org/model_optimization/guide/pruning/

A small usage example

The usage of this API can be as simple as this:

  ```python
  clustering_params = {
    'number_of_clusters': 8,
    'cluster_centroids_init': 'density-based'
  }

  clustered_model = cluster_weights(original_model, **clustering_params)
  ```

This PR does not affect any of the existing pruning functionality and
is completely independent from the rest of the existing code base.

What are the PR files responsible for?

• cluster.py/cluster_test.py - high level clustering interfaces and tests
• cluster_wrapper.py/cluster_wrapper_test.py - an implementation of a
  particular clustering algorithm and associated tests
  clusterable_level.py - similar to PrunableLayer, provides a wrapper for
  a layer to be recognized as a clusterable one.
• clustering_centroids.py/clustering_centroids_test.py - implementations
  of different clusters centroids initialization mechanics and tests for
  them
• clustering_registry.py/clustering_registry_test.py - similar to
  PruningRegistry contains the list of clusterable parameters in each
  layer.

[googlebot]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[googlebot]

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

[akarmi]

I pushed the update that fixes some tests and adds a simple example for clustering MNIST classification convnet.
@alanchiao, can you please review this PR?

[alanchiao]

I started some of comments with "for the future", meaning that I don't expect a code change in the current PR, but some back-and-forth discussion would be nice.

As discussed, I'll follow up with the infra needed to eventually accept this PR.

In the mean time, we'll be able to bring things back into the public API once it's ready for public use (more experiments, docs/tutorials, etc.).

[alanchiao]

Note for self (you can ignore for now): internally we have checks to require that all deps are included directly (e.g. in this fashion, now simplified after this commit).

Going to remind myself what the benefits are and then see if there's an easier way to contribute such changes if still needed.

[akarmi]

@alanchiao, thanks for the feedback. I pushed the update to address your immediate requests. Please let me know if you think we need anything else to merge this.

[alanchiao]

Made additional comments, but I'll try to get things merged now since the PR is large enough the delaying it further is not ideal. We can revisit the comments for future PRs and when finalizing the API.

As I mentioned in one of the comments, I'm pushing to get convergence tests setup (e.g. mobilenet v1 classification for pruning) that can be then used for clustering to make results reproducible. That, together with other experimental result data, is a top priority.

[alanchiao]

For the future: is .register_new_implementation intended to be the way for the end-user to use a custom AbstractClusteringAlgorithm? I'm assuming it is given that the test is in cluster_test.py instead of the registry_test.py.

For pruning, the intended end-user API would be the ClusterableLayer (something different for Quantization due to Quantization specific reasons of worrying about cross-layer interactions). The wrapper can then find the algorithm and use it. The benefit is that for custom Keras layers like this example, the user is already exposed to ClusterableLayer, so we can simplify the API by hiding the registry and register_new_implementation from them.

The downside here is if they would want to use a different algorithm for a built-in Keras layer, then they would have to subclass the built-in Keras layer (which is ok) and then use that custom Keras layer everywhere instead of just calling register_new_implementation once (which is bad).

The choice favors creating the easiest experience for the most common (and basic) users, which really need us to give them sensible defaults for built-in keras layers that work for the majority of the use cases. These users are the focus on tfmot.

For more advanced users and researchers, the experience will be a bit more difficult. Going forward though, we can provide a separate interface where they can effectively register their own registry, to handle the built-in Keras layer with a different algorithm use case.

[alanchiao]

For the future: for both pruning and clustering, we mention the *Registry, but this is something that end-users don't know about how to act on. Only mentioning ClusterableLayer is not great either.

Something clearer would be like: "Layers should be instances of either custom tf.keras layers that subclass ClusterableLayer, or a built-in tf.keras layer supported by the library. For unsupported built-in tf.keras layers, please [TBD next steps ... e.g. create a custom tf.keras layer that subclasses a built-in layer + ClusterableLayer as a workaround and file a Github issue].

[alanchiao]

It'd be nice to pull this end-to-end test into a separate file. That file would be a place we can point users (and others model optimization people) on how far major usage patterns have automatic testing.

[alanchiao]

For the future: would the algorithm still be correct (including the effects on the gradient/backprop) on doing something similar to pruning where we keep the original variables (e.g. self.kernel) and then rely on assignment to update it for the clustering algorithm (what we do by creating the separate 0/1 mask, multiplying by the weights, and then doing an assign)?

I ask because I think doing so would remove the need for gone_variables by removing the need for setattr.

We'll be setting up some convergence tests for pruning (I'm working with another team to make this as accessible as possible, including initial integration), which is something you can eventually utilize to help verify this kind of algorithmic change, in addition to any added e2e testing.

[alanchiao]

Notes on changes I made to get internal tests to pass:

1. changed import of tf.keras.layer.Layer to use public API, which is needed at head due to this commit on keras/layers/init.py
2. Modified BUILD files to include dependencies on tensorflow in unit tests and having py_test use PYTHON_VERSION="PY3"

You'll also need to run pylint, as suggested in this PR and fix the various issues or ask pylint to ignore them.

Otherwise, 1) may be a continuing pain point going forward until the move to the tf. public API usage. That said, I don't think it needs to be addressed until more e2e testing and convergence tests exist.

[alanchiao]

Internally, I can make changes for tests to run on both python 2 and 3. That said, we'll be eventually deprecating Python 2 in #211, though the timing is TBD. Perhaps sometime after quantization-aware training is launched and we make this more visible in the docs.