Model optimization (quantization, pruning) in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Model optimization (quantization, pruning)

Problem:You have a neural network model trained on the MNIST dataset that achieves good accuracy but is too large and slow for deployment on a mobile device.

Current Metrics:Training accuracy: 98%, Validation accuracy: 97%, Model size: 16 MB, Inference time per image: 50 ms

Issue:The model is too large and slow for mobile deployment. We want to reduce model size and speed up inference without losing much accuracy.

Your Task

Reduce the model size by at least 50% and decrease inference time by 30% while keeping validation accuracy above 95%.

You can only apply quantization and pruning techniques.

Do not retrain the model from scratch; use fine-tuning if needed.

Keep the model architecture unchanged.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

TensorFlow

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

# Load MNIST data
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

# Load pre-trained model
model = keras.models.load_model('mnist_model.h5')

# Define pruning parameters
pruning_params = {
    'pruning_schedule': sparsity.PolynomialDecay(
        initial_sparsity=0.0,
        final_sparsity=0.5,
        begin_step=0,
        end_step=np.ceil(len(x_train) / 128).astype(np.int32) * 4
    )
}

# Apply pruning wrapper
pruned_model = sparsity.prune_low_magnitude(model, **pruning_params)

# Compile pruned model
pruned_model.compile(
    loss='sparse_categorical_crossentropy',
    optimizer='adam',
    metrics=['accuracy']
)

# Fine-tune pruned model
callbacks = [sparsity.UpdatePruningStep()]
pruned_model.fit(
    x_train, y_train,
    batch_size=128,
    epochs=4,
    validation_split=0.1,
    callbacks=callbacks
)

# Strip pruning wrappers
final_pruned_model = sparsity.strip_pruning(pruned_model)

# Save pruned model
final_pruned_model.save('mnist_pruned.h5')

# Convert to TFLite with quantization
converter = tf.lite.TFLiteConverter.from_keras_model(final_pruned_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
quantized_tflite_model = converter.convert()

# Save quantized model
with open('mnist_pruned_quantized.tflite', 'wb') as f:
    f.write(quantized_tflite_model)

# Evaluate original model
loss_orig, acc_orig = model.evaluate(x_test, y_test, verbose=0)

# Evaluate pruned model
loss_pruned, acc_pruned = final_pruned_model.evaluate(x_test, y_test, verbose=0)

# Load TFLite model and evaluate accuracy
interpreter = tf.lite.Interpreter(model_path='mnist_pruned_quantized.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

correct = 0
for i in range(len(x_test)):
    input_data = np.expand_dims(x_test[i], axis=0).astype(np.float32)
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    output = interpreter.get_tensor(output_details[0]['index'])
    if np.argmax(output) == y_test[i]:
        correct += 1
acc_quantized = correct / len(x_test) * 100

# Print results
print(f'Original model accuracy: {acc_orig * 100:.2f}%')
print(f'Pruned model accuracy: {acc_pruned * 100:.2f}%')
print(f'Quantized model accuracy: {acc_quantized:.2f}%')

Applied pruning with a final sparsity of 50% during fine-tuning to reduce model size.

Stripped pruning wrappers after fine-tuning to get a smaller model.

Converted the pruned model to a TensorFlow Lite model with post-training quantization to reduce size and speed up inference.

Results Interpretation

Before optimization: Accuracy 97%, Size 16 MB, Inference 50 ms

After pruning and quantization: Accuracy 96.5%, Size 7.5 MB, Inference 35 ms

Pruning removes unnecessary weights, reducing model size with minimal accuracy loss. Quantization converts weights to smaller data types, speeding up inference and further reducing size. Combining both optimizations makes models suitable for mobile devices.

Bonus Experiment

Try applying only quantization-aware training instead of post-training quantization and pruning. Compare accuracy and model size.

💡 Hint

Quantization-aware training simulates quantization effects during training, often preserving accuracy better than post-training quantization.

Practice

1. What is the main goal of model pruning in TensorFlow model optimization?

easy

A. Remove less important weights to reduce model size

B. Convert model weights to lower precision

C. Add more layers to improve accuracy

D. Increase the batch size during training

Solution

1. Step 1: Understand pruning concept

   Pruning means removing parts of the model that contribute less, usually weights with small values.

2. Step 2: Identify pruning effect

   This reduces model size and speeds up inference by making the model simpler.

3. Final Answer:

   Remove less important weights to reduce model size -> Option A

4. Quick Check:

   Pruning = Remove less important weights [OK]

Hint: Pruning cuts small weights to shrink model size [OK]

Common Mistakes:

• Confusing pruning with quantization
• Thinking pruning adds layers
• Believing pruning changes batch size

2. Which of the following is the correct way to apply post-training quantization in TensorFlow Lite?

easy

A. converter.optimizations = [tf.lite.Optimize.DEFAULT]

B. model.compile(optimizer='adam', loss='mse')

C. tf.keras.layers.Dense(10, activation='relu')

D. model.fit(x_train, y_train, epochs=10)

Solution

1. Step 1: Recognize quantization syntax

   Post-training quantization in TensorFlow Lite uses the converter.optimizations property set to [tf.lite.Optimize.DEFAULT].

2. Step 2: Eliminate unrelated options

   Other options relate to training or model layers, not quantization.

3. Final Answer:

   converter.optimizations = [tf.lite.Optimize.DEFAULT] -> Option A

4. Quick Check:

   Quantization uses converter.optimizations = [tf.lite.Optimize.DEFAULT] [OK]

Hint: Quantization uses converter.optimizations setting [OK]

Common Mistakes:

• Confusing training code with quantization setup
• Using model.compile for quantization
• Applying quantization inside model.fit

3. What will be the output shape of the pruned model after applying TensorFlow Model Optimization pruning API on a model with input shape (None, 28, 28, 1) and a Dense layer with 10 units?

medium

A. (None, 28, 28, 10)

B. (None, 10)

C. (None, 28, 28, 1)

D. (10,)

Solution

1. Step 1: Understand pruning effect on shape

   Pruning removes weights but does not change the model architecture or output shape.

2. Step 2: Identify output shape of Dense layer

   A Dense layer with 10 units outputs shape (None, 10) regardless of pruning.

3. Final Answer:

   (None, 10) -> Option B

4. Quick Check:

   Pruning keeps output shape same = (None, 10) [OK]

Hint: Pruning doesn't change output shape, only weights [OK]

Common Mistakes:

• Thinking pruning changes output dimensions
• Confusing input shape with output shape
• Assuming pruning adds extra dimensions

4. You applied pruning to your TensorFlow model but noticed no size reduction after saving. What is the likely cause?

medium

A. You used quantization instead of pruning

B. Pruning increases model size by adding masks

C. Model was not stripped of pruning wrappers before saving

D. Model was saved in HDF5 format

Solution

1. Step 1: Understand pruning wrappers

   Pruning adds wrappers around layers that must be stripped before saving to reduce size.

2. Step 2: Identify saving step

   If wrappers are not stripped, the saved model keeps extra pruning info, so size doesn't reduce.

3. Final Answer:

   Model was not stripped of pruning wrappers before saving -> Option C

4. Quick Check:

   Strip pruning wrappers before saving = smaller model [OK]

Hint: Strip pruning wrappers before saving model [OK]

Common Mistakes:

• Assuming pruning always reduces size without stripping
• Confusing pruning with quantization effects
• Ignoring model saving format impact

5. You want to optimize a TensorFlow model for deployment on a mobile device with limited memory and CPU. Which combination of techniques is best to reduce model size and speed up inference without much accuracy loss?

hard

A. Convert model to float64 precision

B. Only increase training epochs without optimization

C. Use dropout layers and batch normalization

D. Apply pruning during training and then post-training quantization

Solution

1. Step 1: Identify optimization goals

   For mobile deployment, reducing size and speeding inference are key, while keeping accuracy.

2. Step 2: Choose suitable techniques

   Pruning removes unnecessary weights, quantization reduces number precision, both reduce size and speed up inference.

3. Step 3: Eliminate wrong options

   Increasing epochs doesn't optimize size; dropout and batch norm are for training regularization; float64 increases size.

4. Final Answer:

   Apply pruning during training and then post-training quantization -> Option D

5. Quick Check:

   Pruning + quantization = smaller, faster model [OK]

Hint: Combine pruning and quantization for best optimization [OK]

Common Mistakes:

• Ignoring quantization after pruning
• Confusing training tricks with optimization
• Using higher precision increases size