Download notebook

ONNX Quantization

ONNX2TFLite provides two packages. The one used for quantizing an ONNX model is called ONNX2Quant. Quantization requires a calibration dataset to determine the value ranges within the model.

To run quantization, you need to provide input ONNX model and calibration dataset.

1. Setup and configuration

from pathlib import Path
import requests
import zipfile

2. Prepare your model

Let’s prepare the model for quantization. If you already have your own model, you can use that; otherwise, download the example model provided.

model_path = Path("your_model.onnx")
example_model_url = "https://eiq.nxp.com/training-materials/_misc/models/model.onnx"

with open(model_path, "wb") as f:
    response = requests.get(
        url=example_model_url
    )
    f.write(response.content)

    if response.status_code != 200:
        print(f"Failed to download model: {response.content}")
    else:
        print(f"Model downloaded and stored in {model_path}")

3. Prepare your calibration dataset

Calibration dataset has to have specific format and structure. Each model input should have a separate folder containing the input data files. Each file must be a numpy array saved as a .npy file. Following script demonstrates how to prepare calibration dataset for model with single FLOAT32 input:

from pathlib import Path
import numpy as np
from PIL import Image

# Configuration
SOURCE_DIR = Path("images")
OUTPUT_DIR = Path("calibration_data")
OUTPUT_DIR.mkdir(exist_ok=True)

# Convert images to .npy files
for idx, img_path in enumerate(SOURCE_DIR.glob("*.*")):
    img = Image.open(img_path).resize((224, 224))
    img_array = np.array(img, dtype=np.float32) / 255.0
    np.save(OUTPUT_DIR / f"{idx:04d}.npy", img_array)

In case you don’t have your calibration dataset prepared yet, you can use predefined one.

dataset_path = Path("your_calibration_dataset.zip")

example_dataset_url = "https://eiq.nxp.com/training-materials/_misc/datasets/kws_calib.zip"

with open(dataset_path, "wb") as f:
    response = requests.get(
        url=example_dataset_url
    )
    f.write(response.content)

    if response.status_code != 200:
        print(f"Failed to download model: {response.content}")
    else:
        output_dataset_dir = Path("unzipped_dataset")

        # unzip dataset
        with zipfile.ZipFile(dataset_path) as zip_file:
            zip_file.extractall(output_dataset_dir)

        print(f"Dataset downloaded and unzipped in {output_dataset_dir} folder.")

4. Run Quantization

Now we can use the onnx2quant tool to quantize the model using the calibration dataset. The calibration dataset is passed using the -c argument. You need to provide a dataset for each input tensor in the format:

<input>;<dataset_dir>

If the model has multiple inputs, you can repeat the -c argument multiple times. For example, if the model has inputs x and y, you would specify their calibration datasets like this:

-c "x;<path_to_x_dataset>" -c "y;<path_to_y_dataset>"

The example model has only one input, so we only need to set this argument once. If you’re using your own model, make sure to adjust these settings accordingly. Once quantization completes, the quantized model will appear in the current directory under the name <input_model_name>_quant.onnx

input_calib_dataset_directory = output_dataset_dir / "images"
!onnx2quant $model_path -c "images;$input_calib_dataset_directory" -o output_model_quant.onnx

print(f"Quantized model saved to output_model_quant.onnx file")

5. Verify Quantization Results

quantized_model = Path("output_model_quant.onnx")

original_size = model_path.stat().st_size / (1024 * 1024)  # MB
quantized_size = quantized_model.stat().st_size / (1024 * 1024)  # MB
reduction = (1 - quantized_size / original_size) * 100

print(f"Original:  {original_size:.2f} MB")
print(f"Quantized: {quantized_size:.2f} MB")
print(f"Reduction: {reduction:.1f}%")

6. Additional resources

To list all configurable options, use the help command:

!onnx2quant -h