Are we able to prune a pre-trained model? Example: MobileNetV2

user_12540447 0
Viewed 1687

I'm trying to prune a pre-trained model: MobileNetV2 and I got this error. Tried searching online and couldn't understand. I'm running on Google Colab.

These are my imports.

import tensorflow as tf
import tensorflow_model_optimization as tfmot
import tensorflow_datasets as tfds
from tensorflow import keras

import os
import numpy as np
import matplotlib.pyplot as plt
import tempfile
import zipfile

This is my code.

model_1 = keras.Sequential([
    basemodel,
    keras.layers.GlobalAveragePooling2D(),
    keras.layers.Dense(1)                            
])

model_1.compile(optimizer='adam',
                loss=keras.losses.BinaryCrossentropy(from_logits=True),
                metrics=['accuracy'])

model_1.fit(train_batches,
            epochs=5,
            validation_data=valid_batches)

prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

pruning_params = {
    'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=0.50,
                                                             final_sparsity=0.80,
                                                             begin_step=0,
                                                             end_step=end_step)
}


model_2 = prune_low_magnitude(model_1, **pruning_params)

model_2.compile(optmizer='adam',
                loss=keres.losses.BinaryCrossentropy(from_logits=True),
                metrics=['accuracy'])

This is the error i get.

---> 12 model_2 = prune_low_magnitude(model, **pruning_params)

ValueError: Please initialize `Prune` with a supported layer. Layers should either be a `PrunableLayer` instance, or should be supported by the PruneRegistry. You passed: <class 'tensorflow.python.keras.engine.training.Model'>
pruning keras tensorflow python
5 Answers
user_57381500

I faced the same issue with:

  • tensorflow version: 2.2.0

Just updating the version of tensorflow to 2.3.0 solved the issue, I think Tensorflow added support to this feature in 2.3.0.

user_130713400

One thing I found is that the experimental preprocessing I added to my model was throwing this error. I had this at the beginning of my model to help add some more training samples but the keras pruning code doesn't like subclassed models like this. Similarly, the code doesn't like the experimental preprocessing like I have with centering of the image. Removing the preprocessing from the model solved the issue for me.

def classificationModel(trainImgs, testImgs):
  L2_lambda = 0.01
  data_augmentation = tf.keras.Sequential(
  [ layers.experimental.preprocessing.RandomFlip("horizontal", input_shape=IM_DIMS),
    layers.experimental.preprocessing.RandomRotation(0.1),
    layers.experimental.preprocessing.RandomZoom(0.1),])

  model = tf.keras.Sequential()
  model.add(data_augmentation)
  model.add(layers.experimental.preprocessing.Rescaling(1./255, input_shape=IM_DIMS))
...
user_90584960

Saving the model as below and reloading worked for me.

_, keras_file = tempfile.mkstemp('.h5')
tf.keras.models.save_model(model, keras_file, include_optimizer=False)
print('Saved baseline model to:', keras_file)
user_143130680

Had the same problem today, its the following error.

If you don't want the layer to be pruned or don't care for it, you can use this code to only prune the prunable layers in a model:

from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
from tensorflow_model_optimization.python.core.sparsity.keras import prune_registry

def apply_pruning_to_prunable_layers(layer):
    if isinstance(layer, prunable_layer.PrunableLayer) or hasattr(layer, 'get_prunable_weights') or prune_registry.PruneRegistry.supports(layer):
        return tfmot.sparsity.keras.prune_low_magnitude(layer)
    print("Not Prunable: ", layer)
    return layer

model_for_pruning = tf.keras.models.clone_model(
    base_model,
    clone_function=apply_pruning_to_pruneable_layers
)
user_200192180
Related