Shrinking TensorFlow Models: Post-Training Quantization, QAT, and Pruning

Every model you ship has two costs: the one-time cost of training it and the recurring cost of running it. Once a model is in production, inference is the bill that never stops — GPU hours, cold-start latency, memory on a phone, battery on a sensor. Model optimization is how you cut that bill without retraining from scratch, and on the TensorFlow stack the tools are mature and cheap to try: post-training quantization, quantization-aware training (QAT), and pruning.

This tutorial is the piece our other guides deliberately skip. Serving TensorFlow models in production makes a model fast to call; migrating to LiteRT makes it run on-device. This one makes the model itself smaller and cheaper before either of those steps.

Decide what you are optimizing for

Pick one number before you touch any code, because the techniques pull in different directions:

• Model size on disk / over the wire — matters for mobile app bundles and edge OTA updates. Weight-only quantization wins here.
• Inference latency — matters for interactive products. Full integer quantization on a CPU or NPU wins here.
• Memory footprint — matters for microcontrollers and crowded GPU boxes.
• Throughput per dollar — matters for batch scoring. Often dynamic batching (a serving concern) beats quantization.

And decide your accuracy budget up front: "no more than 1 % absolute drop in top-1 on the held-out set" is a decision you can test against. Without it, every optimization discussion turns into opinion.

Build the measurement harness first

You cannot optimize what you are not measuring. Before any conversion, capture a baseline on a fixed evaluation set and a fixed machine:

import time, numpy as np, keras

model = keras.saving.load_model("model.keras")

def benchmark(predict_fn, samples, warmup=10, runs=100):
    for x in samples[:warmup]:
        predict_fn(x)
    t0 = time.perf_counter()
    for i in range(runs):
        predict_fn(samples[i % len(samples)])
    return (time.perf_counter() - t0) / runs * 1000  # ms/inference

baseline_ms = benchmark(lambda x: model(x[None, ...], training=False), eval_images)
baseline_acc = model.evaluate(eval_ds, verbose=0)[1]
print(f"baseline: {baseline_ms:.2f} ms, acc {baseline_acc:.4f}")

Record the file size too (du -h). Every number below is meaningless without this row to compare against.

1. Post-training quantization: try this first

Post-training quantization (PTQ) needs no retraining and takes minutes. It runs through the LiteRT/TFLite converter, and it has three flavours worth knowing.

Dynamic range quantization — weights to int8, activations quantized on the fly. Roughly 4x smaller, meaningful CPU speedup, and usually the smallest accuracy hit. It is the default starting point:

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_dynamic = converter.convert()
open("model_dynamic.tflite", "wb").write(tflite_dynamic)

Float16 quantization — halves size, keeps accuracy essentially intact, and is a good fit for GPU delegates:

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
open("model_fp16.tflite", "wb").write(converter.convert())

Full integer quantization — weights and activations in int8, which is what NPUs, DSPs and microcontrollers actually require. This one needs a representative dataset so the converter can calibrate activation ranges:

def representative_data_gen():
    for images, _ in calibration_ds.batch(1).take(200):
        yield [tf.cast(images, tf.float32)]

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
open("model_int8.tflite", "wb").write(converter.convert())

Two things go wrong here more than anything else. First, the representative dataset must look like production traffic — a couple of hundred real, un-augmented samples covering the range of inputs you expect. Calibrating on synthetic or pre-normalized-differently data is the single most common cause of "quantization destroyed my accuracy". Second, setting inference_input_type = tf.int8 changes your API: callers now send quantized integers, and your client code has to apply the scale and zero point from the interpreter's input details.

Score every artifact through the interpreter, not the Keras model:

from ai_edge_litert.interpreter import Interpreter

interp = Interpreter(model_path="model_int8.tflite")
interp.allocate_tensors()
inp, out = interp.get_input_details()[0], interp.get_output_details()[0]
scale, zero_point = inp["quantization"]

correct = 0
for x, y in eval_pairs:
    q = np.clip(x / scale + zero_point, -128, 127).astype(np.int8)
    interp.set_tensor(inp["index"], q[None, ...])
    interp.invoke()
    correct += int(np.argmax(interp.get_tensor(out["index"])) == y)
print("int8 accuracy:", correct / len(eval_pairs))

2. Quantization-aware training: when PTQ costs too much accuracy

If int8 PTQ blows your accuracy budget — common with small models, heavy depthwise convolutions, or wide activation ranges — QAT simulates quantization during training so the weights adapt to it. It costs a short fine-tune, typically a few epochs at a low learning rate, and usually recovers most of the gap.

import tensorflow_model_optimization as tfmot

q_model = tfmot.quantization.keras.quantize_model(model)
q_model.compile(
    optimizer=keras.optimizers.Adam(1e-5),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)
q_model.fit(train_ds, validation_data=val_ds, epochs=3)

converter = tf.lite.TFLiteConverter.from_keras_model(q_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
open("model_qat_int8.tflite", "wb").write(converter.convert())

Notes from real engagements:

• Start QAT from your trained weights, never from scratch. It is a fine-tune, not a training run.
• Use a learning rate roughly 10–100x lower than the original training run.
• Not every layer type has a quantization wrapper. Use quantize_annotate_layer to quantize only the parts you can, and expect to leave a few layers in float.
• The TensorFlow Model Optimization Toolkit's Keras API tracks the classic tf.keras path. If you are on Keras 3 with a non-TensorFlow backend, do the QAT step on the TensorFlow backend — the same split we describe in the Keras 3 multi-backend guide.

3. Pruning and weight clustering: size wins, speed maybe

Magnitude pruning zeroes out the least important weights over a schedule during fine-tuning:

prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude
schedule = tfmot.sparsity.keras.PolynomialDecay(
    initial_sparsity=0.20, final_sparsity=0.75,
    begin_step=0, end_step=len(train_ds) * 3,
)

pruned = prune_low_magnitude(model, pruning_schedule=schedule)
pruned.compile(optimizer=keras.optimizers.Adam(1e-5),
               loss="sparse_categorical_crossentropy", metrics=["accuracy"])
pruned.fit(train_ds, epochs=3,
           callbacks=[tfmot.sparsity.keras.UpdatePruningStep()])

final = tfmot.sparsity.keras.strip_pruning(pruned)   # do not skip this

Be clear-eyed about the payoff. Unstructured sparsity compresses beautifully — a 75 %-sparse model gzips to a fraction of its size — but on ordinary CPU and GPU kernels it does not make inference faster, because the zeros are still multiplied. You get speed only with hardware or runtimes that exploit sparsity, or with structured sparsity patterns. Treat pruning as a download-size and memory technique, and stack it with quantization (strip_pruning then quantize) rather than choosing between them.

Weight clustering (tfmot.clustering.keras.cluster_weights) is a similar trade: excellent compressibility, no free latency.

4. Server-side options that are not LiteRT

The techniques above are aimed at the LiteRT/edge path, but most of our clients run on servers. There the levers are different:

• XLA compilation via jit_compile=True on model.compile(), or tf.function(jit_compile=True) for the serving graph. Kernel fusion often buys 10–30 % on GPU for free.
• Mixed precision — keras.mixed_precision.set_global_policy("mixed_float16") (or mixed_bfloat16) is usually a bigger and safer win on modern GPUs than any int8 scheme.
• Dynamic batching in TF Serving, covered in the serving tutorial. For throughput-bound workloads this beats quantization, and it costs no accuracy at all.
• Distillation — train a small student on the large model's outputs. More work than everything above, but the only technique that reliably gives an order-of-magnitude reduction.

Try those before you spend a sprint on int8 for a server deployment.

5. Record the results in a table you can defend

Every optimization project we run ends in the same artifact:

Variant	Size	p50 latency	p99 latency	Accuracy	Δ vs baseline
float32 baseline	44 MB	38 ms	61 ms	0.912	—
float16	22 MB	34 ms	55 ms	0.912	−0.000
dynamic range int8	11 MB	19 ms	31 ms	0.907	−0.005
full int8 (PTQ)	11 MB	12 ms	20 ms	0.881	−0.031
full int8 (QAT)	11 MB	12 ms	21 ms	0.906	−0.006
pruned 75 % + int8	4 MB gz	12 ms	21 ms	0.903	−0.009

Those numbers are illustrative, not a promise — the shape of the table is the point. Measure latency on the target device, not your laptop, and report p99 alongside p50, because tail latency is what users feel. With a table like this, choosing a variant becomes a business decision instead of an argument.

Common failure modes

• Calibration data that does not match production. Fix the representative dataset before you blame the technique.
• Comparing accuracy on the training set. Always use a held-out set the optimization never touched.
• Optimizing a model that was never the bottleneck. Profile end to end first; if 80 % of your request time is feature lookup, a 4x faster model changes nothing.
• No provenance. Store the conversion script, calibration dataset hash and toolkit versions next to the artifact. A .tflite file nobody can reproduce is technical debt.
• Silent input-scale drift. Once inputs are int8, a client that forgets the scale and zero point produces plausible-looking garbage. Pin it in a contract test.

Where to start

Take your current production model, run dynamic range quantization, and measure. That is an afternoon of work and it frequently delivers a 4x size reduction and a solid CPU speedup for a rounding error of accuracy. Only if that is not enough do you escalate to full int8, QAT, and pruning — in that order, measuring at every step.

If you would like help sizing that work, our edge AI and LiteRT deployment practice does exactly this, and you can get in touch to talk through your model and your latency target.