tensorflow/model-optimization

Support for tensorflow hub layers

Opened this issue · 1 comments

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.

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.