How to prune an existing tensorflow/keras model trained on imagenet

Question

I am trying to prune InceptionNetV3 from keras trained on imagenet, right now I am using a tensorflow-datasets which has a subset of imagenet which I use for pruning. Currently my pruned models do not work and returns garbage data when tested using the same dataset it was pruned on. How do I prune without losing all accuracy? Here is my code:

Imports:

import logging
import tempfile
from pathlib import Path
import tensorflow as tf
from tensorflow import keras
import numpy as np
import tensorflow_datasets as tfds
from cv2 import cv2  # Pylint now views cv2 as a library
import tensorflow_model_optimization as tfmot

All of these imports are up to date, I'm currently using Python 3.10.1.

Here is the code I am using to prune the model.

v2_path = 'C:\\temp\\imagenet_v2'

inception_image_size = (299, 299)

image_count = 5

batch_size = 512

epochs = 4

dataset = tfds.load(name='imagenet_v2', split='test', data_dir=v2_path)
numpy_dataset = tfds.as_numpy(dataset)

layer_count = 313

count = [1]


def main():
    v2_full_path = 'C:\\temp\\imagenet_v2\\downloads\\extracted\\TAR_GZ.s3-us-west-2_image_image-match-frequ8MN_35JZFrGeoTI82aIgjNtpWbosMu7yp_w5ODXJynw.tar.gz\\imagenetv2-matched-frequency-format-val'

    dataset_train = tf.keras.utils.image_dataset_from_directory(directory=v2_full_path, image_size=inception_image_size,
                                                                label_mode='categorical')

    inception_model = tf.keras.applications.InceptionV3(weights='imagenet',
                                                        pooling='avg',
                                                        input_shape=(299, 299, 3))

    def apply_pruning_to_dense(layer):
        count[0] += 1  # Python throws a fit if I use a normal variable, but doesn't mind layer_count

        if layer_count - count[0] < 5:
            return tfmot.sparsity.keras.prune_low_magnitude(layer)
        return layer

    model_for_pruning = tf.keras.models.clone_model(
        inception_model,
        clone_function=apply_pruning_to_dense,
    )

    inception_model = tf.keras.applications.InceptionV3(weights="imagenet")

    logdir = tempfile.mkdtemp()

    callbacks = [
        tfmot.sparsity.keras.UpdatePruningStep(),
        tfmot.sparsity.keras.PruningSummaries(log_dir=logdir),
    ]

    model_for_pruning.compile(loss='categorical_crossentropy',
                              optimizer=keras.optimizers.SGD(learning_rate=1e-3),
                              metrics=['accuracy'])

    model_for_pruning.fit(dataset_train,
                          batch_size=batch_size,
                          epochs=epochs,
                          callbacks=callbacks,
                          use_multiprocessing=True)

    save_test_model(inception_model, ".tflite")

    save_test_model(model_for_pruning, "_prune.tflite")

When I run the model through model_for_pruning.fit(...) the accuracy rating is only around 1% - 2%. Though it used to be around .16% per epoch. I fixed this by adding label_mode='categorical' when obtaining the dataset which leads me to believe that the issue is somehow with either my dataset or how I use it.

The resulting pruned tensorflow lite model has a 0% accuracy rating when tested against the imagenet_v2 subset while the upruned one gets around a 40% accuracy rating.

How to prune an existing tensorflow/keras model trained on imagenet

0 Answers0