This release is a pre-release and may not be stable for production use.
AI Edge Quantizer
A quantizer for advanced developers to quantize converted LiteRT models. It aims to facilitate advanced users to strive for optimal performance on resource demanding models (e.g., GenAI models).
Build Status
| Build Type | Status |
|---|---|
| Unit Tests (Linux) |  |
| Nightly Release |  |
| Nightly Colab |  |
Installation
Requirements and Dependencies
- Python versions: 3.10, 3.11, 3.12, 3.13
- Operating system: Linux, MacOS
- TensorFlow:
Install
Nightly PyPi package:
pip install ai-edge-quantizer-nightly
API Usage
The quantizer requires two inputs:
- An unquantized source LiteRT model (with FP32 data type in the FlatBuffer
format with
.tfliteextension) - A quantization recipe (details below)
and outputs a quantized LiteRT model that's ready for deployment on edge devices.
Basic Usage
In a nutshell, the quantizer works according to the following steps:
- Instantiate a
Quantizerclass. This is the entry point to the quantizer's functionalities that the user accesses. - Load a desired quantization recipe (details in subsection).
- Quantize (and save) the model. This is where most of the quantizer's internal logic works.
from ai_edge_quantizer import quantizer, recipe
qt = quantizer.Quantizer("path/to/input/tflite")
# Load a ready-to-use recipe (for example dynamic int8 quantization).
qt.load_quantization_recipe(recipe.dynamic_wi8_afp32())
qt.quantize().export_model("/path/to/output/tflite")
Please see the getting started colab for the simplest quick start guide on those 3 steps, and the selective quantization colab with more details on advanced features.
LiteRT Model
Please refer to the LiteRT documentation for ways to generate LiteRT models from Jax, PyTorch and TensorFlow. The input source model should be an FP32 (unquantized) model in the FlatBuffer format with .tflite extension.
Quantization Recipe
The user needs to specify a quantization recipe using AI Edge Quantizer's API to apply to the source model. The quantization recipe encodes all information on how a model is to be quantized, such as number of bits, data type, symmetry, scope name, etc.
Essentially, a quantization recipe is defined as a collection of commands of the following type:
“Apply Quantization Algorithm X on Operator Y under Scope Z with ConfigN”.
For example:
"Uniformly quantize the FullyConnected op under scope 'dense1/' with INT8 symmetric with Dynamic Quantization".
All the unspecified ops will be kept as FP32 (unquantized). The scope of an operator in TFLite is defined as the output tensor name of the op, which preserves the hierarchical model information from the source model (e.g., scope in TF). The best way to obtain scope name is by visualizing the model with Model Explorer.
Currently, there are three ways to quantize an operator:
-
dynamic quantization (recommended): weights are quantized while activations remain in a float format and are not processed by AI Edge Quantizer (AEQ). The runtime kernel handles the on-the-fly quantization of these activations, as identified by
compute_precision=integerandexplicit_dequantize=False.- Pros: reduced model size and memory usage. Latency improvement due to integer computation. No sample data requirement (calibration).
- Cons: on-the-fly quantization of activation tensors can affect model quality. Not supported in all hardware (e.g., some GPU and NPU).
-
weight only quantization: only model weights are quantized, not activations. The actual operation (op) computation remains in float. The quantized weight is explicitly dequantized before being fed into the op, by inserting a dequantize op between the quantized weight and the consuming op. To enable this,
compute_precisionwill be set tofloatandexplicit_dequantizetoTrue.- Pros: reduced model size and memory usage. No sample data requirement (calibration). Usually has the best model quality.
- Cons: no latency benefit (may be worse) due to float computation with explicit dequantization.
-
static quantization: both weights and activations are quantized. This requires a calibration phase to estimate quantization parameters of runtime tensors (activations).
- Pros: reduced model size, memory usage, and latency.
- Cons: requires sample data for calibration. Imposing static quantization parameters (derived from calibration) on runtime tensors can compromise quality.
Generally, we recommend dynamic quantization for CPU/GPU deployment and static quantization for NPU deployment.
We include commonly used recipes in recipe.py. This is demonstrated in the getting started colab example. Advanced users can build their own recipe through the quantizer API.
Model Validation & Accuracy Benchmarking
Quantizing a model inherently introduces numerical noise. After calling
qt.quantize(), you can verify the mathematical distortion between the float
baseline and the quantized model using the built-in validate() method, which
returns a single ComparisonResult object mapping nodes to their error metric
values. You can print them or automatically save them to Model Explorer JSON
files:
# 1. Default validation (evaluates MSE metric by default)
comparison_results = qt.validate(test_data=sample_data)
print(
"Per-layer metrics:",
comparison_results.get_all_tensor_results(),
)
# 2. Multi-metric validation (save all metrics and validation json data directly)
comparison_results = qt.validate(
test_data=sample_data,
error_metrics=[
quantizer.ValidationErrorMetric.MSE,
quantizer.ValidationErrorMetric.SNR,
],
save_folder='/tmp/'
)
all_results = comparison_results.get_all_tensor_results()
for tensor_name, metrics in all_results.items():
print(
f"Tensor: {tensor_name} "
f"- MSE: {metrics.get(quantizer.ValidationErrorMetric.MSE.value, 0.0):.6f} "
f"- SNR: {metrics.get(quantizer.ValidationErrorMetric.SNR.value, 0.0):.6f}"
)
More detailed examples can be found in quantize_toy_model.py.
Visualizing Models with Model Explorer
The best way to obtain exact operator scope names and visually compare tensor shapes and quantization scales between baseline float and quantized graphs is using Model Explorer.
To visualize two exported .tflite models side-by-side in your terminal, run:
model_explorer --models \
"/path/to/baseline_float.tflite,/path/to/quantized_model.tflite"
Deployment
Please refer to the LiteRT deployment documentation for ways to deploy a quantized LiteRT model.
Advanced Recipes
There are many ways the user can configure and customize the quantization recipe beyond using a template in recipe.py. For example, the user can configure the recipe to achieve these features:
- Selective quantization (exclude selected ops from being quantized)
- Flexible mixed scheme quantization (mixture of different precision, compute precision, scope, op, config, etc)
- 4-bit weight quantization
The selective quantization colab shows some of these more advanced features.
For specifics of the recipe schema, please refer to the OpQuantizationRecipe
in [recipe_manager.py].
For advanced usage involving mixed quantization, the following API may be useful:
- Use
Quantizer:load_quantization_recipe()in quantizer.py to load a custom recipe. - Use
Quantizer:update_quantization_recipe()in quantizer.py to extend or override specific parts of the recipe.
Operator coverage
The table below outlines the allowed configurations for available recipes.
| Config | DYNAMIC_WI8_AFP32 | DYNAMIC_WI4_AFP32 | DYNAMIC_WI4_AFP32_BLOCKWISE | DYNAMIC_WI2_AFP32_BLOCKWISE | STATIC_WI8_AI8 | STATIC_WI8_AI16 | STATIC_WI4_AI8 | STATIC_WI4_AI16 | WEIGHTONLY_WI8_AFP32 | WEIGHTONLY_WI4_AFP32 | |
| activation | num_bits | None | None | None | None | 8 | 16 | 8 | 16 | None | None |
| symmetric | None | None | None | None | [TRUE, FALSE] | TRUE | [TRUE, FALSE] | TRUE | None | None | |
| granularity | None | None | None | None | TENSORWISE | TENSORWISE | TENSORWISE | TENSORWISE | None | None | |
| dtype | None | None | None | None | INT | INT | INT | INT | None | None | |
| weight | num_bits | 8 | 4 | 4 | 2 | 8 | 8 | 4 | 4 | 8 | 4 |
| symmetric | TRUE | TRUE | TRUE | TRUE | TRUE | TRUE | TRUE | TRUE | [TRUE, FALSE] | [TRUE, FALSE] | |
| granularity | [CHANNELWISE, TENSORWISE] | [CHANNELWISE, TENSORWISE] | [BLOCKWISE_32, BLOCKWISE_64, BLOCKWISE_128, BLOCKWISE_256] | [BLOCKWISE_32, BLOCKWISE_64, BLOCKWISE_128, BLOCKWISE_256] | [CHANNELWISE, TENSORWISE] | [CHANNELWISE, TENSORWISE] | [CHANNELWISE, TENSORWISE] | [CHANNELWISE, TENSORWISE] | [CHANNELWISE, TENSORWISE] | [CHANNELWISE, TENSORWISE] | |
| dtype | INT | INT | INT | INT | INT | INT | INT | INT | INT | INT | |
| explicit_dequantize | FALSE | FALSE | FALSE | FALSE | FALSE | FALSE | FALSE | FALSE | TRUE | TRUE | |
| compute_precision | INTEGER | INTEGER | INTEGER | INTEGER | INTEGER | INTEGER | INTEGER | INTEGER | FLOAT | FLOAT |
Quantization Support for Operators with Weights
| Config | DYNAMIC_WI8_AFP32 | DYNAMIC_WI4_AFP32 | DYNAMIC_WI4_AFP32_BLOCKWISE | DYNAMIC_WI2_AFP32_BLOCKWISE | STATIC_WI8_AI8 | STATIC_WI8_AI16 | STATIC_WI4_AI8 | STATIC_WI4_AI16 | WEIGHTONLY_WI8_AFP32 | WEIGHTONLY_WI4_AFP32 |
| BATCH_MATMUL | ✓ |
✓ |
✓ |
✓ |
✓ |
|||||
| CONV_2D | ✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
||
| CONV_2D_TRANSPOSE | ✓ |
✓ |
✓ |
✓ |
||||||
| DEPTHWISE_CONV_2D | ✓ |
✓ |
✓ |
✓ |
||||||
| EMBEDDING_LOOKUP | ✓ |
✓ |
✓ |
✓ |
✓ |
|||||
| FULLY_CONNECTED | ✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
Quantization Support for Activations-Only Operators
| Config | STATIC_WI8_AI8 | STATIC_WI8_AI16 |
| ADD | ✓ |
✓ |
| AVERAGE_POOL_2D | ✓ |
✓ |
| BROADCAST_TO | ✓ |
✓ |
| CONCATENATION | ✓ |
✓ |
| DIV | ✓ |
✓ |
| DYNAMIC_UPDATE_SLICE | ✓ |
✓ |
| EQUAL | ✓ |
✓ |
| GATHER | ✓ |
✓ |
| GATHER_ND | ✓ |
✓ |
| GELU | ✓ |
✓ |
| HARD_SWISH | ✓ |
|
| LOGISTIC | ✓ |
✓ |
| MAX_POOL_2D | ✓ |
✓ |
| MAXIMUM | ✓ |
✓ |
| MEAN | ✓ |
✓ |
| MIRROR_PAD | ✓ |
✓ |
| MUL | ✓ |
✓ |
| NOT_EQUAL | ✓ |
✓ |
| PACK | ✓ |
✓ |
| PAD | ✓ |
✓ |
| PADV2 | ✓ |
✓ |
| REDUCE_MIN | ✓ |
✓ |
| RELU | ✓ |
✓ |
| RESHAPE | ✓ |
✓ |
| RESIZE_BILINEAR | ✓ |
✓ |
| RESIZE_NEAREST_NEIGHBOR | ✓ |
✓ |
| RSQRT | ✓ |
✓ |
| SELECT | ✓ |
✓ |
| SELECT_V2 | ✓ |
✓ |
| SLICE | ✓ |
✓ |
| SOFTMAX | ✓ |
✓ |
| SPACE_TO_DEPTH | ✓ |
|
| SPLIT | ✓ |
✓ |
| SQRT | ✓ |
✓ |
| SQUARED_DIFFERENCE | ✓ |
|
| STRIDED_SLICE | ✓ |
✓ |
| SUB | ✓ |
✓ |
| SUM | ✓ |
✓ |
| TANH | ✓ |
✓ |
| TRANSPOSE | ✓ |
✓ |
| UNPACK | ✓ |
✓ |
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file ai_edge_quantizer_nightly-0.9.0.dev20260820.tar.gz.
File metadata
- Download URL: ai_edge_quantizer_nightly-0.9.0.dev20260820.tar.gz
- Upload date:
- Size: 263.4 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.12.5 {"installer":{"name":"uv","version":"0.12.5","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
1838a706484a34f3d2eb24cbee423cdfac2ad515d02d37816bf341d031b793ad
|
|
| MD5 |
7a3dd8b8c99e4cbc96cfe0a166de20f3
|
|
| BLAKE2b-256 |
9882d5f2b98a9c826d3648215ad6340eb02cbec9eef1317923dfc37f515db5f9
|
File details
Details for the file ai_edge_quantizer_nightly-0.9.0.dev20260820-py3-none-any.whl.
File metadata
- Download URL: ai_edge_quantizer_nightly-0.9.0.dev20260820-py3-none-any.whl
- Upload date:
- Size: 480.0 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.12.5 {"installer":{"name":"uv","version":"0.12.5","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
46f5954ad0c8d637a4ca7c08fed6e8c5bb423b0763d5e704a4a6b43ff4557432
|
|
| MD5 |
367a307b722e0959b8e32defa43b8271
|
|
| BLAKE2b-256 |
5fadaf8aa85c83aff69c5ac001a74d6c722875a9880ad906e7941902d161ab15
|
