tensorflow_keras_magnitude_weight_pruning_mnist_quickstart.py

1import tensorflow as tf
2import numpy as np
3import tensorflow_model_optimization as tfmot
4
5# 1. Load and prepare the dataset
6mnist = tf.keras.datasets.mnist
7(train_images, train_labels), (test_images, test_labels) = mnist.load_data()
8
9# Normalize the input image so that each pixel value is between 0 and 1.
10train_images = train_images / 255.0
11test_images = test_images / 255.0
12
13# 2. Define the baseline model
14model = tf.keras.Sequential([
15  tf.keras.layers.InputLayer(input_shape=(28, 28)),
16  tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
17  tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),
18  tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
19  tf.keras.layers.Flatten(),
20  tf.keras.layers.Dense(10)
21])
22
23# 3. Apply pruning to the model
24prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude
25
26# Compute end step to finish pruning after 2 epochs.
27batch_size = 128
28epochs = 2
29validation_split = 0.1 # 10% of training set will be used for validation set. 
30
31num_images = train_images.shape[0] * (1 - validation_split)
32end_step = np.ceil(num_images / batch_size).astype(np.int32) * epochs
33
34# Define parameters for pruning.
35pruning_params = {
36      'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=0.50,
37                                                               final_sparsity=0.80,
38                                                               begin_step=0,
39                                                               end_step=end_step)
40}
41
42model_for_pruning = prune_low_magnitude(model, **pruning_params)
43
44# 4. Recompile and train the model
45model_for_pruning.compile(optimizer='adam',
46              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
47              metrics=['accuracy'])
48
49callbacks = [
50  tfmot.sparsity.keras.UpdatePruningStep(),
51]
52
53model_for_pruning.fit(train_images, train_labels,
54                  batch_size=batch_size, epochs=epochs, validation_split=validation_split,
55                  callbacks=callbacks)
56
57# 5. Export the model for compression
58# Strip the pruning wrappers to keep only the pruned weights
59model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)
60
61# Evaluate the final model
62_, model_for_pruning_accuracy = model_for_pruning.evaluate(
63   test_images, test_labels, verbose=0)
64
65print('Pruned test accuracy:', model_for_pruning_accuracy)