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

Question

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'>

Answer 1

I believe you are following Pruning in Keras Example and jumped into Fine-tune pre-trained model with pruning section without setting your prunable layers. You have to reinstantiate model and set layers you wish to set as prunable. Follow this guide for further information on how to set prunable layers.

https://www.tensorflow.org/model_optimization/guide/pruning/comprehensive_guide.md