Quantizing a Model for TFLite

Quantizing a model for TFLite takes about six lines. Knowing whether the result is still the same model takes the rest of this page, and skipping it is how a quantized model reaches production with quietly worse predictions.

Which quantization you need

The converter offers several modes and they are not interchangeable. Pick by target, not by size:

• Dynamic range. Weights to int8, activations computed in float at runtime. Roughly a 4x size reduction, no calibration data needed. Fine for CPU, and not accepted by integer-only accelerators.
• Float16. Halves the size, keeps float arithmetic. Useful on GPU delegates, useless for an integer NPU.
• Full integer. Every tensor int8 or uint8, including activations, which requires calibration data. This is the only mode an Edge TPU or a comparable integer-only NPU will accept, and it is what the rest of this page does.

The reason full integer needs data and the others do not is worth holding on to. Weights are known at conversion time, so their range can be read straight off the tensor. Activations only exist while the model is running, so their range has to be observed — which means running the model on real inputs.

There is a second axis worth knowing before you start, because it explains why weights survive quantization better than activations do. Weight quantization is normally per-channel: each output channel of a convolution gets its own scale, so one channel with an unusually wide distribution does not drag the rest down. Activation quantization is per-tensor: one scale for the whole tensor, because the values do not exist until runtime and there is nowhere to attach a per-channel scale. That asymmetry is why the failures described at the end of this page are almost always activation failures.

The representative dataset

A generator that yields batches of inputs in the same distribution and the same preprocessing as production. Google’s documentation uses around 100 samples and yields one at a time:

import tensorflow as tf

def representative_data_gen():
    for input_value in tf.data.Dataset.from_tensor_slices(train_images).batch(1).take(100):
        yield [input_value]

Three things go wrong here and all three are silent:

• Wrong preprocessing. If production normalises to [-1, 1] and your generator yields [0, 255], every activation scale is calibrated against the wrong range and the model is ruined in a way no exception reports.
• Unrepresentative samples. 100 images of one class calibrate for one class. Sample across the real distribution, including whatever your edge cases are.
• Too few, or too many. A handful of samples miss the tails; thousands mostly cost conversion time. A few hundred spanning the distribution is the documented shape.

Converting

1. Configure the converter for full integer quantization. The target_spec.supported_ops line is the one that makes this strict: with TFLITE_BUILTINS_INT8 alone, an operator with no integer implementation raises an error at conversion instead of silently leaving a float island in the graph.

   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.uint8
   converter.inference_output_type = tf.uint8
   
   tflite_quant_model = converter.convert()
   open("model_quant.tflite", "wb").write(tflite_quant_model)

2. Decide the input and output types deliberately. Setting inference_input_type and inference_output_type to tf.uint8 or tf.int8 removes the float conversion layers at the model boundary, which integer-only accelerators require. Leave them at float and you get a model with quantize/dequantize ops at each end — convenient in Python, rejected by some hardware.
3. Check the sizes. The int8 file should land near a quarter of the float32 one. If it is close to half, you produced a float16 model; if it barely moved, the optimization flag did not take.

   ls -l model_float.tflite model_quant.tflite

4. Inspect the resulting tensor types rather than trusting the flags. This prints the dtype the interpreter actually sees at the boundary:

   interpreter = tf.lite.Interpreter(model_path="model_quant.tflite")
   interpreter.allocate_tensors()
   print(interpreter.get_input_details()[0]["dtype"])
   print(interpreter.get_output_details()[0]["dtype"])
   print(interpreter.get_input_details()[0]["quantization"])

   The quantization field gives the (scale, zero_point) pair. A scale of 0.0 means that tensor was not quantized.

Verifying against the float model

This is the step that makes the conversion trustworthy, and it is the one most tutorials leave out. Run the same held-out inputs through both models and compare, remembering to dequantize the integer output before comparing it to a float one:

import numpy as np

def run_tflite(path, images, quantized):
    interp = tf.lite.Interpreter(model_path=path)
    interp.allocate_tensors()
    inp = interp.get_input_details()[0]
    out = interp.get_output_details()[0]
    preds = []
    for image in images:
        x = np.expand_dims(image, axis=0)
        if quantized:
            scale, zero_point = inp["quantization"]
            x = (x / scale + zero_point).astype(inp["dtype"])
        interp.set_tensor(inp["index"], x)
        interp.invoke()
        y = interp.get_tensor(out["index"])[0]
        if quantized:
            scale, zero_point = out["quantization"]
            y = (y.astype(np.float32) - zero_point) * scale
        preds.append(y)
    return np.array(preds)

float_preds = run_tflite("model_float.tflite", test_images, quantized=False)
quant_preds = run_tflite("model_quant.tflite", test_images, quantized=True)

agree = np.mean(np.argmax(float_preds, 1) == np.argmax(quant_preds, 1))
print("top-1 agreement:", agree)
print("max abs output delta:", np.abs(float_preds - quant_preds).max())

Two numbers, and they answer different questions. Agreement is how often the two models pick the same class, which is what your application experiences. Maximum absolute delta is how far the logits moved, which tells you how close the model is to disagreeing on the next input you have not tested. High agreement with a large delta is a model on the edge; check it on more data before shipping.

Compare against a held-out set, not against the representative dataset. The quantized model was calibrated on those samples, so measuring on them flatters it in exactly the way you are trying to detect.

When accuracy collapses

A model that loses a point or two of accuracy is normal. One that falls off a cliff has a specific cause, and it is almost always dynamic range:

• An outlier activation. One layer whose activations occasionally reach a value hundreds of times the typical magnitude forces a huge scale factor on that tensor, which quantizes every ordinary value to the same handful of integers. Per-channel weight quantization — the default for convolutions — does not help, because the problem is in the activations.
• Wide-range operations. Softmax, exponentials, layer normalisation and anything with a division span ranges int8 cannot hold well. Where the tooling supports it, leaving these in float and quantizing the rest is the standard escape.
• Calibration mismatch. Covered above, and the first thing to re-check because it is the cheapest to fix.

If post-training quantization cannot be made to hold, the documented next step is quantization-aware training: the model learns with the quantization error present, so it adapts to the coarser numeric grid rather than being projected onto it afterwards. It costs a training run, which is why it is second. The mechanism behind all of this — why outlier features specifically are what break low-bit formats — is the same one described in quantization at inference time.