TFLite Quantizer to AI Edge Quantizer Migration Guide

This guide helps you migrate from legacy post-training quantization in the TensorFlow Lite Converter (TFLQ) to AI Edge Quantizer (AEQ).

Key Changes: The 2-Step Workflow

In legacy TFLQ, quantization was tightly coupled into the model conversion step inside tf.lite.TFLiteConverter. In AEQ, model conversion and quantization are decoupled into two distinct, framework-agnostic steps:

  1. Step 1: Export unquantized model. Convert your source model (from PyTorch, JAX, Keras, or TensorFlow) into an unquantized (FP32) .tflite flatbuffer.
  2. Step 2: Quantize with AEQ. Run AI Edge Quantizer using the standalone aeq command-line tool or the ai_edge_quantizer Python API.

Dynamic Range Quantization (INT8 Weights & FP32 Activations)

Dynamic range quantization quantizes weights to 8-bit integers while activations remain in floating point and are dynamically quantized at runtime during inference.

Legacy TFLQ

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()

Modern AEQ

Command Line (aeq):

aeq --model_file="model.tflite" \
    --recipe=dynamic_wi8_afp32 \
    --output_dir="/path/to/output"

Python API:

from ai_edge_quantizer import quantizer

qt = quantizer.Quantizer("model.tflite")
qt.load_quantization_recipe("dynamic_wi8_afp32")
quantized_model = qt.quantize()

Full Integer Static Quantization (INT8 Weights & INT8 Activations)

Full integer static quantization quantizes both weights and activations into 8-bit integers using representative calibration data to calculate quantization parameters.

Option A: Integer with Float Fallback (Default Float Input / Output)

In legacy TFLQ, float fallback quantized internal weights and activations to integer while keeping the graph's input and output tensors in float32 for compatibility with standard application pipelines. If an operation lacked a quantized implementation, it also fell back to float.

Legacy TFLQ (post_training_integer_quant):

import tensorflow as tf

def representative_dataset():
  for data in dataset:
    yield {
      "image": data.image.astype(np.float32),
      "bias": data.bias.astype(np.float32),
    }

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
tflite_quant_model = converter.convert()

Modern AEQ:

from ai_edge_quantizer import algorithm_manager, calibrator, quantizer, qtyping
import numpy as np

qt = quantizer.Quantizer("model.tflite")

qt.load_quantization_recipe("static_wi8_ai8")
# Disable quantization for input and output tensors.
qt.update_quantization_recipe(
    regex='.*',
    operation_name=qtyping.TFLOperationName.INPUT,
    algorithm_key=algorithm_manager.AlgorithmName.NO_QUANTIZE,
)
qt.update_quantization_recipe(
    regex='.*',
    operation_name=qtyping.TFLOperationName.OUTPUT,
    algorithm_key=algorithm_manager.AlgorithmName.NO_QUANTIZE,
)

calibration_data = {
    "serving_default": [
        {"image": data.image.astype(np.float32),
         "bias": data.bias.astype(np.float32)}
        for data in dataset
    ]
}
calibration_result = qt.calibrate(
    calibration_data,
    mode=calibrator.CalibrationMode.CALIBRATION_PROFILER_BASED,
)
quantized_model = qt.quantize(calibration_result=calibration_result)

Option B: Integer Only (Integer Input / Output)

For integer-only hardware, microcontrollers, and accelerators (e.g., NPUs or Edge TPU) that don't support floating-point arithmetic, the model input and output tensors must also be quantized to integer.

Legacy TFLQ (post_training_integer_quant):

import tensorflow as tf

def representative_dataset():
  for data in dataset:
    yield {
      "image": data.image.astype(np.float32),
      "bias": data.bias.astype(np.float32),
    }

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8  # or tf.uint8
converter.inference_output_type = tf.int8  # or tf.uint8
tflite_quant_model = converter.convert()

Modern AEQ:

from ai_edge_quantizer import calibrator, quantizer
import numpy as np

qt = quantizer.Quantizer("model.tflite")
qt.load_quantization_recipe("static_wi8_ai8")
calibration_data = {
    "serving_default": [
        {"input_tensor_name": sample.astype(np.float32)}
        for sample in calibration_samples
    ]
}
calibration_result = qt.calibrate(
    calibration_data,
    mode=calibrator.CalibrationMode.CALIBRATION_PROFILER_BASED,
)
quantized_model = qt.quantize(calibration_result=calibration_result)

Integer Quantization with 16-bit Activations (INT8 Weights & INT16 Activations)

Quantizes weights to 8-bit integers and activations to 16-bit integers. This scheme provides higher precision than INT8 activations while retaining smaller INT8 weights.

Legacy TFLQ

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
]
converter.representative_dataset = representative_data_gen
tflite_quant_model = converter.convert()

Modern AEQ

from ai_edge_quantizer import calibrator, quantizer
import numpy as np

qt = quantizer.Quantizer("model.tflite")
qt.load_quantization_recipe("static_wi8_ai16")

calibration_data = {
    "serving_default": [
        {"input_tensor_name": sample.astype(np.float32)}
        for sample in calibration_samples
    ]
}
calibration_result = qt.calibrate(
    calibration_data,
    mode=calibrator.CalibrationMode.CALIBRATION_PROFILER_BASED,
)
quantized_model = qt.quantize(calibration_result=calibration_result)

Float16 Weights Quantization

Float16 weights quantization (non-uniform float casting) converts model weights to 16-bit floating-point numbers without quantizing activations.

Advantages:

  • Reduces model size by up to half, since all weights become half of their original size
  • Minimal loss in accuracy
  • Some delegates (e.g. the GPU delegate) can operate directly on float16 data, resulting in faster execution than float32 computations

Disadvantages:

  • Latency may not improve as much as with fixed-point quantization schemes
  • By default on CPU, a float16 quantized model will "dequantize" the weights values to float32 during inference

Legacy TFLQ

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
tflite_quant_model = converter.convert()

Modern AEQ

In AEQ, float16 weights quantization is supported through non-uniform float casting:

from ai_edge_quantizer import algorithm_manager, quantizer

qt = quantizer.Quantizer("model.tflite")
qt.add_weight_only_config(
    regex=".*",
    operation_name="*",
    num_bits=16,
    algorithm_key=algorithm_manager.AlgorithmName.FLOAT_CASTING,
)
quantized_model = qt.quantize()

Recipe Cheat Sheet

Quantization Scheme Target Hardware Calibration? AEQ Recipe / Config
Dynamic Range INT8 CPU / GPU No "dynamic_wi8_afp32"
Dynamic Range INT4 CPU / GPU No "dynamic_wi4_afp32"
Static INT8 / INT8 NPU Yes "static_wi8_ai8"
Static INT8 / INT16 NPU Yes "static_wi8_ai16"