Quantization-Aware Training and Pruning

When post-training quantization isn't accurate enough

Post-training quantization is convenient but can hurt accuracy on sensitive tasks. Quantization-Aware Training (QAT) simulates int8 quantization during training's forward pass — the model stays float32, but learns weights robust to int8 conversion. A final TFLite conversion produces the actual int8 model.

Typical accuracy gain over post-training quantization: 1–3% absolute. Significant when you're at the edge of your accuracy budget.

Magnitude-based pruning zeros out small weights, creating sparse models. Combined with quantization, can shrink models ~10× with minimal accuracy loss. The TF Model Optimization Toolkit (tensorflow_model_optimization) provides Keras layer wrappers that gradually increase sparsity during training.

Code

QAT 워크플로우·python

import tensorflow as tf
import tensorflow_model_optimization as tfmot

# pip install tensorflow-model-optimization

# 1. Start with a pretrained float32 model
base_model = tf.keras.applications.MobileNetV2(
    weights='imagenet', input_shape=(224, 224, 3),
)
# ... assume base_model is fine-tuned on your task ...

# 2. Wrap with QAT — every layer gets fake quant ops
q_aware = tfmot.quantization.keras.quantize_model(base_model)

# 3. Recompile (wrapper changes layer structure)
q_aware.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy'],
)

# 4. Fine-tune for a few epochs — learns quantization-robust weights
q_aware.fit(train_data, train_labels, batch_size=64, epochs=5,
            validation_split=0.1)

# 5. Convert to actual int8 TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(q_aware)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_qat = converter.convert()

with open('model_qat.tflite', 'wb') as f:
    f.write(tflite_qat)

Magnitude pruning·python

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

prune = tfmot.sparsity.keras.prune_low_magnitude

# Schedule: ramp 50% → 80% sparsity over training
end_step = np.ceil(num_images / batch_size).astype(int) * epochs
pruning_params = {
    'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
        initial_sparsity=0.50,
        final_sparsity=0.80,
        begin_step=0,
        end_step=end_step,
    )
}

pruned = prune(model, **pruning_params)
pruned.compile(optimizer='adam',
               loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
               metrics=['accuracy'])

# REQUIRED callback — without it, sparsity doesn't update
callbacks = [
    tfmot.sparsity.keras.UpdatePruningStep(),
    tfmot.sparsity.keras.PruningSummaries(log_dir='/tmp/pruning'),
]
pruned.fit(x_train, y_train, batch_size=batch_size, epochs=epochs,
           callbacks=callbacks)

# Strip wrappers before export
final = tfmot.sparsity.keras.strip_pruning(pruned)
converter = tf.lite.TFLiteConverter.from_keras_model(final)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
pruned_quant_tflite = converter.convert()
# Result: typically ~10× smaller than original float32 model

Quantization-Aware Training and Pruning

When post-training quantization isn't accurate enough

Code

Progress

Comments 0