Deep neural networks often suffer from large parameter sizes and lengthy training times, making deployment challenging in resource-constrained environments. Network pruning offers a solution by reducing the number of parameters without significantly impacting accuracy.

Unlike traditional approaches that focused on weight removal, recent research suggests pruning is more akin to model architecture search, primarily removing filters or neurons rather than individual weights. This structural change requires removing corresponding feature maps and weights, fundamentally altering the network architecture.

The pruning process follows an Oracle Pruning methodology: identifying and removing unimportant parameters, then retraining to restore model performance. The effectiveness of this approach depends heavily on the criteria used to evaluate parameter importance.

Implementation Steps

1. Train a Keras model to acceptable accuracy levels
2. Prepare the model for pruning
3. Create a pruning schedule with multiple training iterations
4. Export the pruned model by removing pruning wrappers
5. Convert the Keras model to TensorFlow Lite (optional)

Installation Requirements

Begin by installing the necessary packages:

pip uninstall -yq tensorflow
pip uninstall -yq tf-nightly
pip install -Uq tf-nightly-gpu
pip install -q tensorflow-model-optimization

Model Preparation

Load your pre-trained model and convert it to a "prunable" format. The pruning algorithm will be applied to all layers that support weight pruning. For our pruning schedule, we'll start with 50% sparsity and gradually increase to 90% over training iterations.

import numpy as np
import tensorflow as tf
from tensorflow_model_optimization.sparsity import keras as sparsity

# Load your pre-trained model
model = tf.keras.models.load_model('pretrained_model.h5')

# Define training parameters
epochs = 4
end_step = np.ceil(1.0 * num_training_samples / batch_size).astype(np.int32) * epochs

# Configure pruning parameters
pruning_params = {
    'pruning_schedule': sparsity.PolynomialDecay(
        initial_sparsity=0.50,
        final_sparsity=0.90,
        begin_step=0,
        end_step=end_step,
        frequency=100)
}

# Apply pruning to the model
pruned_model = sparsity.prune_low_magnitude(model, pruning_params)
pruned_model.compile(
    loss=tf.keras.losses.categorical_crossentropy,
    optimizer='adam',
    metrics=['accuracy'])

Model Training with Pruning

Train the pruned model with appropriate callbacks to track the pruning process:

# Set up callbacks for pruning
callbacks = [
    sparsity.UpdatePruningStep(),
    sparsity.PruningSummaries(log_dir='./logs', profile_batch=0)]

# Train the model
pruned_model.fit(training_data, training_labels,
                batch_size=batch_size,
                epochs=epochs,
                callbacks=callbacks,
                validation_data=(validation_data, validation_labels))

# Evaluate the model
test_loss, test_accuracy = pruned_model.evaluate(test_data, test_labels)
print(f'Test loss: {test_loss}')
print(f'Test accuracy: {test_accuracy}')

Finalizing the Pruned Model

After training, remove the pruning wrappers to finalize the model:

final_model = sparsity.strip_pruning(pruned_model)
final_model.summary()

Verifying Sparsity

You can verify the percentage of zero weights in each layer:

for i, weights in enumerate(final_model.get_weights()):
    print(f"{final_model.weights[i].name} -- Total: {weights.size}, Zeros: {np.sum(weights == 0) / weights.size * 100:.2f}%")

Model Compression

The pruned model can now be compressed significantly. For example, a model that was originally 12.52 MB might be reduced to 2.51 MB using standard compression techniques:

import tempfile
import zipfile
import os

# Save the pruned model
_, pruned_model_path = tempfile.mkstemp(".h5")
tf.keras.models.save_model(final_model, pruned_model_path, include_optimizer=False)

# Compress the model
_, compressed_path = tempfile.mkstemp(".zip")
with zipfile.ZipFile(compressed_path, "w", compression=zipfile.ZIP_DEFLATED) as f:
    f.write(pruned_model_path)

print(f"Size of pruned model before compression: {os.path.getsize(pruned_model_path) / float(2**20):.2f} MB")
print(f"Size of pruned model after compression: {os.path.getsize(compressed_path) / float(2**20):.2f} MB")

Conversion to TensorFlow Lite

For deployment on mobile devices, convert the pruned model to TensorFlow Lite format:

# Create TFLite model
converter = tf.lite.TFLiteConverter.from_keras_model(final_model)
converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
tflite_model = converter.convert()

# Save the TFLite model
tflite_model_path = "/tmp/pruned_model.tflite"
with open(tflite_model_path, "wb") as f:
    f.write(tflite_model)

Evaluating the TensorFlow Lite Model

You can evaluate the accuracy of the converted model as follows:

def evaluate_tflite_model(interpreter, test_data, test_labels):
    input_index = interpreter.get_input_details()[0]["index"]
    output_index = interpreter.get_output_details()[0]["index"]
    
    correct_predictions = 0
    total_samples = 0
    
    for sample, label in zip(test_data, test_labels):
        input_data = sample.reshape((1, *sample.shape))
        interpreter.set_tensor(input_index, input_data)
        interpreter.invoke()
        output = interpreter.get_tensor(output_index)
        
        if np.argmax(output) == np.argmax(label):
            correct_predictions += 1
        total_samples += 1
    
    return correct_predictions / total_samples

# Load and evaluate the TFLite model
interpreter = tf.lite.Interpreter(model_path=tflite_model_path)
interpreter.allocate_tensors()
tflite_accuracy = evaluate_tflite_model(interpreter, test_data, test_labels)
print(f"TFLite model accuracy: {tflite_accuracy}")