Model optimization for serving (quantization, pruning) in MLOps - Commands & Configuration

Introduction

Model optimization helps make machine learning models smaller and faster. This is important when you want to run models on devices with limited power or speed. Techniques like quantization and pruning reduce model size and improve serving speed without losing much accuracy.

When you want to deploy a model on a mobile phone with limited memory and CPU power.

When you need faster predictions from a model in a web service to handle more users.

When you want to reduce cloud costs by using smaller models that need less compute.

When you want to run models on edge devices like IoT sensors with low resources.

When you want to improve battery life on devices running AI models by reducing computation.

Commands

This Python code loads a pretrained MobileNetV2 model, applies quantization to reduce its size and improve speed, then saves the optimized model as a TensorFlow Lite file for efficient serving.

Terminal

import tensorflow as tf
from tensorflow import keras

# Load a pretrained model
model = keras.applications.MobileNetV2(weights='imagenet')

# Convert to a TensorFlow Lite model with quantization
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()

# Save the quantized model to disk
with open('mobilenetv2_quant.tflite', 'wb') as f:
    f.write(tflite_quant_model)

print('Quantized model saved as mobilenetv2_quant.tflite')

Expected Output

Quantized model saved as mobilenetv2_quant.tflite

This code applies pruning to the MobileNetV2 model to remove less important weights, making the model smaller. It runs a short training loop to apply pruning, then saves the pruned model for serving.

Terminal

import tensorflow_model_optimization as tfmot
from tensorflow import keras

# Load the original model
model = keras.applications.MobileNetV2(weights='imagenet')

# Apply pruning to the model
prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude
pruning_params = {
    'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
        initial_sparsity=0.0, final_sparsity=0.5, begin_step=0, end_step=1000)
}

model_for_pruning = prune_low_magnitude(model, **pruning_params)

# Compile the pruned model
model_for_pruning.compile(optimizer='adam', loss='categorical_crossentropy')

# Dummy training loop to apply pruning
import numpy as np
x_dummy = np.random.rand(10, 224, 224, 3)
y_dummy = np.random.rand(10, 1000)
model_for_pruning.fit(x_dummy, y_dummy, epochs=1, steps_per_epoch=10)

# Strip pruning wrappers to get the final pruned model
model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)

# Save the pruned model
model_for_export.save('mobilenetv2_pruned.h5')

print('Pruned model saved as mobilenetv2_pruned.h5')

Expected Output

Epoch 1/1

10/10 [==============================] - 3s 200ms/step - loss: 6.908

Pruned model saved as mobilenetv2_pruned.h5

This command loads the quantized TensorFlow Lite model and prepares it for inference. It prints the input and output details so you know how to feed data and read predictions.

Terminal

import tensorflow as tf

# Load the quantized TFLite model
interpreter = tf.lite.Interpreter(model_path='mobilenetv2_quant.tflite')
interpreter.allocate_tensors()

# Get input and output details
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

print('Input details:', input_details)
print('Output details:', output_details)

Expected Output

Input details: [{'name': 'input_1', 'index': 0, 'shape': array([ 1, 224, 224, 3], dtype=int32), 'dtype': <class 'numpy.float32'>, 'quantization': (0.0, 0)}]

Output details: [{'name': 'Logits', 'index': 123, 'shape': array([ 1, 1000], dtype=int32), 'dtype': <class 'numpy.float32'>, 'quantization': (0.0, 0)}]

Key Concept

If you remember nothing else from this pattern, remember: quantization and pruning reduce model size and speed up serving by simplifying the model without losing much accuracy.

Code Example

MLOps

import tensorflow as tf
from tensorflow import keras
import tensorflow_model_optimization as tfmot
import numpy as np

# Load pretrained model
model = keras.applications.MobileNetV2(weights='imagenet')

# Quantization
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()
with open('mobilenetv2_quant.tflite', 'wb') as f:
    f.write(tflite_quant_model)
print('Quantized model saved as mobilenetv2_quant.tflite')

# Pruning
prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude
pruning_params = {
    'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
        initial_sparsity=0.0, final_sparsity=0.5, begin_step=0, end_step=1000)
}
model_for_pruning = prune_low_magnitude(model, **pruning_params)
model_for_pruning.compile(optimizer='adam', loss='categorical_crossentropy')
x_dummy = np.random.rand(10, 224, 224, 3)
y_dummy = np.random.rand(10, 1000)
model_for_pruning.fit(x_dummy, y_dummy, epochs=1, steps_per_epoch=10)
model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)
model_for_export.save('mobilenetv2_pruned.h5')
print('Pruned model saved as mobilenetv2_pruned.h5')

# Load and prepare quantized model for inference
interpreter = tf.lite.Interpreter(model_path='mobilenetv2_quant.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print('Input details:', input_details)
print('Output details:', output_details)

Common Mistakes

Skipping the short training step after applying pruning.

Pruning requires training to identify and remove less important weights; skipping this means pruning has no effect.

Run a brief training loop after applying pruning to let the model adjust and remove weights properly.

Trying to use the original TensorFlow model directly after quantization without converting to TensorFlow Lite.

Quantization for serving usually requires conversion to TensorFlow Lite format; the original model won't benefit from quantization optimizations.

Always convert the quantized model to TensorFlow Lite format before deployment.

Not allocating tensors before running inference with the TensorFlow Lite interpreter.

Without allocating tensors, the interpreter cannot prepare memory for inputs and outputs, causing errors.

Call interpreter.allocate_tensors() before inference.

Summary

Use TensorFlow Lite converter with optimization flags to quantize models for smaller size and faster serving.

Apply pruning with TensorFlow Model Optimization Toolkit by wrapping the model and running a short training loop.

Always allocate tensors before running inference on TensorFlow Lite models to prepare input and output buffers.

Practice

1. What is the main goal of quantization in model optimization for serving?

easy

A. Increase the size of the model for better performance

B. Reduce the precision of numbers to make the model smaller and faster

C. Add more neurons to improve accuracy

D. Remove entire layers from the model to simplify it

Solution

1. Step 1: Understand quantization purpose

   Quantization reduces the number precision (like from 32-bit to 8-bit) to save memory and speed up computation.

2. Step 2: Compare options

   Removing layers is pruning, adding neurons increases size, increasing size is opposite of optimization.

3. Final Answer:

   Reduce the precision of numbers to make the model smaller and faster -> Option B

4. Quick Check:

   Quantization = Reduce precision [OK]

Hint: Quantization means lowering number precision to save space [OK]

Common Mistakes:

• Confusing pruning with quantization
• Thinking quantization adds complexity
• Believing quantization increases model size

2. Which of the following is the correct syntax to apply pruning using TensorFlow Model Optimization API in Python?

easy

A. pruned_model = tfmot.sparsity.keras.prune_low_magnitude(model, pruning_schedule=pruning_schedule)

B. pruned_model = tf.prune_low_magnitude(model, schedule=pruning_schedule)

C. pruned_model = tfmot.prune_low_magnitude(model, pruning_schedule=pruning_schedule)

D. pruned_model = tfmot.sparsity.prune_low_magnitude(model, pruning_schedule)

Solution

1. Step 1: Recall TensorFlow pruning API structure

   The pruning function is under tfmot.sparsity.keras and requires the pruning_schedule argument.

2. Step 2: Check syntax correctness

   pruned_model = tfmot.sparsity.keras.prune_low_magnitude(model, pruning_schedule=pruning_schedule) matches the correct full path and argument names. Others miss parts or have wrong argument names.

3. Final Answer:

   pruned_model = tfmot.sparsity.keras.prune_low_magnitude(model, pruning_schedule=pruning_schedule) -> Option A

4. Quick Check:

   Correct pruning syntax = pruned_model = tfmot.sparsity.keras.prune_low_magnitude(model, pruning_schedule=pruning_schedule) [OK]

Hint: TensorFlow pruning is under tfmot.sparsity.keras with pruning_schedule [OK]

Common Mistakes:

• Omitting 'keras' in the API path
• Using wrong argument names
• Calling pruning directly from tf module

3. Given the following PyTorch code snippet for quantization, what will be the output type of the model's weights after applying dynamic quantization?

import torch
import torch.nn as nn

model = nn.Linear(10, 5)
quantized_model = torch.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)
print(type(quantized_model.weight()))

medium

A. TypeError: 'weight' is not callable

B.

C. AttributeError: 'Linear' object has no attribute 'weight'

D.

Solution

1. Step 1: Analyze dynamic quantization effect

   torch.quantization.quantize_dynamic converts nn.Linear to torch.nn.quantized.dynamic.Linear, where weight is a method returning dequantized weights as torch.Tensor.

2. Step 2: Trace the print statement

   quantized_model.weight() succeeds, returning a torch.Tensor (fp32 dequantized), so print(type(...)) outputs <class 'torch.Tensor'>.

3. Final Answer:

   <class 'torch.Tensor'> -> Option D

4. Quick Check:

   Dynamic quant: weight() returns Tensor [OK]

Hint: Dynamic quantization makes weight() callable returning Tensor [OK]

Common Mistakes:

• Thinking weight remains non-callable attribute like original Linear
• Confusing quantized_model type with weight type
• Expecting error on quantized model weight access

4. You tried pruning a TensorFlow model but got an error: AttributeError: module 'tensorflow_model_optimization' has no attribute 'sparsity'. What is the most likely cause?

medium

A. The tensorflow_model_optimization package is not installed

B. You used the wrong pruning schedule argument

C. You forgot to import tensorflow_model_optimization as tfmot

D. Pruning is not supported in TensorFlow

Solution

1. Step 1: Understand the error message

   The error says the module has no attribute 'sparsity', which usually means the package is missing or outdated.

2. Step 2: Check common causes

   If the package is not installed, Python cannot find the 'sparsity' submodule. Importing incorrectly or wrong argument causes different errors.

3. Final Answer:

   The tensorflow_model_optimization package is not installed -> Option A

4. Quick Check:

   Missing package = AttributeError [OK]

Hint: Missing package causes AttributeError on submodules [OK]

Common Mistakes:

• Assuming import alias causes error
• Blaming pruning schedule argument
• Thinking pruning unsupported in TensorFlow

5. You want to optimize a large deep learning model for mobile deployment by combining pruning and quantization. Which sequence of steps is best to minimize model size and maintain accuracy?

hard

A. Apply quantization first, then prune the model to remove weights

B. Train the model with quantization-aware training, then prune after deployment

C. First prune the model to remove unimportant weights, then apply quantization to reduce number precision

D. Only prune the model; quantization is not compatible with pruning

Solution

1. Step 1: Understand pruning and quantization order

   Pruning removes unimportant weights first, reducing model size and complexity.

2. Step 2: Apply quantization after pruning

   Quantization then reduces number precision on the smaller pruned model, further shrinking size and speeding inference.

3. Final Answer:

   First prune the model to remove unimportant weights, then apply quantization to reduce number precision -> Option C

4. Quick Check:

   Prune first, then quantize = First prune the model to remove unimportant weights, then apply quantization to reduce number precision [OK]

Hint: Prune first to shrink, then quantize to compress numbers [OK]

Common Mistakes:

• Quantizing before pruning reduces pruning effectiveness
• Thinking pruning and quantization cannot be combined
• Pruning after deployment is too late