TensorFlow Model Optimisation Toolkit

Quantizing a model after training (as covered in Chapter 6) is the simplest way to shrink it, but it is not always the most accurate way. The TensorFlow Model Optimisation Toolkit (TFMOT) provides techniques that bake efficiency into the training process itself, often recovering accuracy that post-training methods alone would lose.

Introduction to Model Optimisation

Model optimisation exists to solve one core tension: bigger, more accurate models are slower and more expensive to run, while smaller, faster models risk losing accuracy. TFMOT offers three complementary techniques, each attacking the problem from a different angle:

• Quantization-aware training: simulates low-precision arithmetic during training so the model learns to be robust to it.

• Pruning: removes weights that contribute little to the model's output, creating sparse, more compact networks.

• Weight clustering: groups similar weight values together, reducing the number of unique values that need to be stored.

These techniques can be combined; a pruned and quantized model is common in production, each multiplying the savings of the other.

Quantization-Aware Training

Post-training quantization, covered in Chapter 6, converts an already-trained float model into int8 after the fact. Quantization-aware training (QAT) instead inserts fake quantization nodes during training itself, so the model experiences the rounding error of low precision while its weights are still being learned, and can adapt to compensate.

import tensorflow_model_optimization as tfmot

model = tf.keras.Sequential([

tf.keras.layers.Dense(256, activation='relu', input_shape=(784,)),

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

])

# Wrap the model to simulate quantization during training

quantize_model = tfmot.quantization.keras.quantize_model

q_aware_model = quantize_model(model)

q_aware_model.compile(optimizer='adam',

loss='sparse_categorical_crossentropy',

metrics=['accuracy'])

q_aware_model.fit(x_train, y_train, epochs=5, validation_split=0.1)

# Convert the QAT model to a genuinely quantized TFLite model

converter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model)

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

quantized_tflite_model = converter.convert()

QAT typically recovers accuracy that would otherwise be lost to post-training quantization, especially for models that are sensitive to precision loss, such as those with very small layers or tight numeric ranges.

Pruning Neural Networks

Pruning gradually zeroes out the smallest-magnitude weights in a layer during training, on the theory that weights close to zero contribute little to the final output. The result is a sparse model, one with a large fraction of zero weights, that compresses extremely well and, with the right hardware or runtime, can also run faster.

import tensorflow_model_optimization as tfmot

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, # 50% of weights pruned to zero

begin_step=0,

end_step=1000,

)

}

model_for_pruning = prune_low_magnitude(model, **pruning_params)

model_for_pruning.compile(optimizer='adam',

loss='sparse_categorical_crossentropy',

metrics=['accuracy'])

callbacks = [tfmot.sparsity.keras.UpdatePruningStep()]

model_for_pruning.fit(x_train, y_train, epochs=10, callbacks=callbacks)

# Strip pruning wrappers before exporting the final model

final_model = tfmot.sparsity.keras.strip_pruning(model_for_pruning)

Sparsity is introduced gradually, following a schedule, rather than all at once, which gives the remaining weights time to adjust and compensate for the ones being zeroed out.

Weight Clustering

Weight clustering reduces the number of unique values in a layer's weights by grouping them into a small number of clusters and replacing every weight with its cluster's centroid value. Fewer unique values mean the model compresses far better, since repeated values are cheap to encode.

import tensorflow_model_optimization as tfmot

cluster_weights = tfmot.clustering.keras.cluster_weights

CentroidInitialization = tfmot.clustering.keras.CentroidInitialization

clustering_params = {

'number_of_clusters': 16,

'cluster_centroids_init': CentroidInitializationLINEAR,

}

clustered_model = cluster_weights(model, **clustering_params)

clustered_model.compile(optimizer='adam',

loss='sparse_categorical_crossentropy',

metrics=['accuracy'])

clustered_model.fit(x_train, y_train, epochs=3)

final_model = tfmot.clustering.keras.strip_clustering(clustered_model)

Clustering pairs especially well with compression formats used for storage and over-the-air model updates, since fewer unique values translate directly into smaller compressed file sizes.

Performance Benchmarking

Optimisation techniques are only worth applying if you can measure their actual impact. A proper benchmark compares the baseline model against each optimised variant across three dimensions:

import time

interpreter = tf.lite.Interpreter(model_path='model.tflite')

interpreter.allocate_tensors()

input_index = interpreter.get_input_details()[0]['index']

output_index = interpreter.get_output_details()[0]['index']

start = time.time()

for sample in x_test[:100]:

interpreter.set_tensor(input_index, sample.reshape(1, 784).astype('float32'))

interpreter.invoke()

_ = interpreter.get_tensor(output_index)

avg_latency_ms = (time.time() - start) / 100 * 1000

print(f'Average latency: {avg_latency_ms:.2f} ms')

Quick Tip

Always benchmark on the actual target hardware, not your development machine. A model that is faster on a desktop CPU is not guaranteed to be faster on a mobile chip, since they have very different strengths in integer versus floating-point throughput.