Support for tensorflow hub layers

tensorflow / model-optimization

A toolkit to optimize ML models for deployment for Keras and TensorFlow, including quantization and pruning.

https://www.tensorflow.org/model_optimization

Apache License 2.0

1.49k stars 319 forks source link

Support for tensorflow hub layers #924

Open albertnanda opened 2 years ago

albertnanda commented 2 years ago

System information

TensorFlow version (you are using): 2.7
Are you willing to contribute it (Yes/No): Yes

Motivation

Pruning and clustering both do not currently support tensorflow hub models. Ex: ValueError: Please initialize Cluster with a supported layer. Layers should either be a ClusterableLayer instance, or should be supported by the ClusteringRegistry. You passed: <class 'tensorflow_hub.keras_layer.KerasLayer'> Most of the NLP models which can benefit from these techniques are hosted on tf hub, hence, the request.

daverim commented 2 years ago

Hi albertnanda@ thanks for your request -- we are currently actively researching how to support these layers. Unfortunately the hub layer is basically a wrapper for a saved model function which does not permit a search for cluster-able weights. Instead we are trying to find alternatives based on manipulating the saved model. This may take awhile.

If you inspect the model after loading and can determine the weights you would like to cluster: a workaround would be to add a custom layer that extends hub layer and implements get_clusterable_weights and clone the downloaded model to use the custom layer you have implemented.