Pruning: Keras subclassed model increased support

[alanchiao]

Currently the pruning API will throw an error when a subclassed model is passed to it. Users can get around this by diving into the subclassed models and applying pruning to individual Sequential/Functional models and tf.keras layers.

Better support is important for various cases (e.g. Object Detection, BERT examples) and issues such as this one.

We can provide better support for pruning an entire subclassed model.

• Pruning some layers of the model would still require going into the model definition itself, though now you can prune a whole subclassed
  model inside a subclassed model.
• This would only prune variables that live inside a tf.keras.Layer (whether a built-in layer or a custom layer using a the PrunableLayer interface).

Implementation-wise, we can iterate through the layers of a subclassed model (and nested models) and applying pruning to all of them. Replacing a layer in an already created model will be tricky and we'd have to do this without clone_model.

[nutsiepully]

It seems to me we are mixing a few issues. I want to make sure I understand the problem correctly. Please correct me if I'm wrong.

Issue is about handling models recursively - a Keras Model which contains another model, not about subclass models.

If one of the layers in the model is Sequential, or Functional we should ideally traverse that model further and prune all the layers within them.

So, if the user provides the the following.

model1 = tf.keras.Sequential([tf.keras.layers.Dense(3)])
model2 = tf.keras.Sequential([model1, tf.keras.layers.Dense(2)])

we should prune both the Dense layers, and recursively traverse the model. I agree with that.

I am not quite sure how the subclass models play a role here. It is possible a user combines a keras Sequential, Functional and Layer within a subclass model. In that case, the user has to explicitly prune whatever they want.

Since we can't reliably clone a subclass model, it isn't possible for us to apply pruning to the model.