quantize-rs

Neural network quantization toolkit for ONNX models, written in Rust with Python bindings.

quantize-rs converts float32 ONNX models to INT8 or INT4 representation using post-training quantization. It supports weight-only quantization, activation-based calibration, per-channel quantization, and outputs standard ONNX QDQ (DequantizeLinear) graphs compatible with ONNX Runtime.

Features

• INT8 and INT4 quantization -- per-tensor or per-channel
• Activation-based calibration -- runs real inference on calibration data via tract to determine optimal quantization ranges
• Multiple calibration methods -- MinMax, Percentile (99.9th), Entropy (KL divergence), MSE
• ONNX QDQ output format -- quantized models use DequantizeLinear nodes and load directly in ONNX Runtime
• Graph connectivity validation -- verifies that every node input resolves after quantization
• Per-layer selection -- exclude layers by name, set per-layer bit widths, or skip small tensors via min_elements
• CLI -- single-model quantization, batch processing, validation, benchmarking, config-file driven workflows
• Python bindings -- via PyO3; install with pip install quantization-rs
• Typed error handling -- QuantizeError enum at all public API boundaries (no more string-parsing anyhow errors)
• Rust library -- usable as a crate dependency; all public items have doc comments
• Property-based tests -- 15 proptest cases covering quantization round-trips, error bounds, and bit-packing
• Criterion benchmarks -- throughput and per-channel comparison benchmarks in benches/

Installation

Python

pip install quantization-rs

Build from source (requires Rust toolchain):

pip install maturin
git clone https://github.com/AR-Kamal/quantize-rs.git
cd quantize-rs
maturin develop --release --features python

Rust CLI

cargo install quantize-rs

As a library dependency

[dependencies]
quantize-rs = "0.6"

Quick start

Python

import quantize_rs

# Weight-based INT8 quantization
quantize_rs.quantize("model.onnx", "model_int8.onnx", bits=8)

# Activation-based calibration (better accuracy)
quantize_rs.quantize_with_calibration(
    "resnet18.onnx",
    "resnet18_int8.onnx",
    calibration_data="samples.npy",
    method="minmax"
)

# Inspect model metadata
info = quantize_rs.model_info("model.onnx")
print(f"{info.name}: {info.num_nodes} nodes")

See Python API documentation for the full reference.

CLI

# INT8 quantization
quantize-rs quantize model.onnx -o model_int8.onnx

# INT4 with per-channel quantization
quantize-rs quantize model.onnx -o model_int4.onnx --bits 4 --per-channel

# Activation-based calibration
quantize-rs calibrate model.onnx \
    --data calibration.npy \
    -o model_calibrated.onnx \
    --method minmax

# Validate a quantized model (structure, connectivity, numerical sanity)
quantize-rs validate model.onnx model_int8.onnx --detailed

# Compare original vs quantized
quantize-rs benchmark model.onnx model_int8.onnx

# Batch processing
quantize-rs batch *.onnx -o quantized/ --bits 8 --per-channel

# Config-file driven workflow
quantize-rs config quantize.yaml --dry-run

Rust library

use quantize_rs::{OnnxModel, Quantizer, QuantConfig};
use quantize_rs::onnx_utils::graph_builder::QdqWeightInput;

fn main() -> anyhow::Result<()> {
    let mut model = OnnxModel::load("model.onnx")?;
    let weights = model.extract_weights();

    let config = QuantConfig {
        bits: 8,
        per_channel: true,
        calibration_method: None,
        ..Default::default()
    };
    let quantizer = Quantizer::new(config);

    let mut quantized_data = Vec::new();
    for weight in &weights {
        let quantized = quantizer.quantize_tensor(
            &weight.data,
            weight.shape.clone(),
        )?;
        let (scales, zero_points) = quantized.get_all_scales_zero_points();
        let is_per_channel = quantized.is_per_channel();

        quantized_data.push(QdqWeightInput {
            original_name: weight.name.clone(),
            quantized_values: quantized.data(),
            scales,
            zero_points,
            bits: quantized.bits(),
            axis: if is_per_channel { Some(0) } else { None },
        });
    }

    model.save_quantized(&quantized_data, "model_int8.onnx")?;
    Ok(())
}

CLI reference

quantize

quantize-rs quantize <MODEL> [OPTIONS]

Options:
  -o, --output <FILE>       Output path [default: model_quantized.onnx]
  -b, --bits <4|8>          Bit width [default: 8]
      --per-channel         Per-channel quantization
      --exclude <LAYER>     Exclude a layer by name (repeatable)
      --min-elements <N>    Skip tensors with fewer than N elements

calibrate

quantize-rs calibrate <MODEL> --data <DATA> [OPTIONS]

Options:
      --data <FILE>       Calibration data (.npy)
  -o, --output <FILE>     Output path [default: model_calibrated.onnx]
  -b, --bits <4|8>        Bit width [default: 8]
      --per-channel       Per-channel quantization
      --method <METHOD>   minmax | percentile | entropy | mse [default: percentile]

batch

quantize-rs batch <MODELS>... -o <DIR> [OPTIONS]

Options:
  -o, --output <DIR>      Output directory (required)
  -b, --bits <4|8>        Bit width [default: 8]
      --per-channel       Per-channel quantization
      --skip-existing     Skip models that already have output files
      --continue-on-error Do not abort on individual model failures

validate

quantize-rs validate <ORIGINAL> <QUANTIZED> [--detailed]

Checks structure preservation, graph connectivity, weight shapes, and numerical sanity (all-zero detection, constant-value detection). With --detailed, prints per-layer error analysis.

benchmark

quantize-rs benchmark <ORIGINAL> <QUANTIZED>

Compares node counts, weight counts, file sizes, and compression ratios.

info

quantize-rs info <MODEL>

Prints model name, opset version, node count, inputs, and outputs.

config

quantize-rs config <CONFIG_FILE> [--dry-run]

Runs quantization from a YAML or TOML configuration file. Example:

bits: 8
per_channel: true

models:
  - input: models/resnet18.onnx
    output: quantized/resnet18_int8.onnx

  - input: models/mobilenet.onnx
    output: quantized/mobilenet_int8.onnx

batch:
  input_dir: "models/*.onnx"
  output_dir: quantized/
  skip_existing: true

How it works

Quantization

Each float32 weight tensor is mapped to a fixed-point integer representation:

scale = (max - min) / (qmax - qmin)
quantized = round(value / scale) + zero_point
dequantized = (quantized - zero_point) * scale

For INT8, the quantized range is [-128, 127]. For INT4, it is [-8, 7]. INT4 values are bit-packed (two values per byte) in memory for 8x compression, but stored as INT8 in ONNX files (DequantizeLinear requires INT8 input in opsets < 21).

Per-channel quantization

Computes separate scale and zero_point for each output channel (axis 0). This is particularly effective when different channels have vastly different weight distributions, which is common in convolutional layers.

Activation-based calibration

Instead of deriving quantization ranges from weight values alone, calibration runs forward passes on representative input samples and records the actual activation distributions at each layer. The observed ranges produce tighter quantization parameters. Four methods are available:

Output format

Quantized models use the ONNX QDQ pattern. For each quantized weight, the original float32 initializer is replaced with:

• {name}_quantized -- INT8 tensor (same shape)
• {name}_scale -- float32 scalar
• {name}_zp -- INT8 scalar
• A DequantizeLinear node whose output is the original tensor name

Because the DequantizeLinear output carries the original name, all downstream nodes (Conv, MatMul, etc.) remain unchanged. The graph loads and runs in ONNX Runtime without modification.

ONNX Runtime integration

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession("model_int8.onnx")
input_name = session.get_inputs()[0].name
x = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = session.run(None, {input_name: x})

Testing

# Rust tests (90 passing: 63 unit + 12 integration + 15 property-based)
cargo test

# With output
cargo test -- --nocapture

# Integration tests requiring model files on disk
cargo test -- --ignored --nocapture

# Criterion benchmarks
cargo bench

# Python tests (requires maturin develop)
pytest test_python_bindings.py -v

Known limitations

• ONNX input only. PyTorch and TensorFlow models must be exported to ONNX first.
• Per-channel DequantizeLinear writes 1-D scale/zero_point tensors with the axis attribute. ONNX Runtime supports this in opset >= 13.
• INT4 values are stored as INT8 bytes in the ONNX file. True 4-bit packing requires opset 21 or a custom operator.
• Activation calibration uses tract for inference; tract may not support all ONNX ops found in large production models.

Contributing

1. Fork the repository
2. Create a feature branch
3. Add tests for new functionality
4. Ensure cargo test and cargo clippy pass
5. Submit a pull request

License

MIT

Acknowledgments

• tract -- ONNX inference engine used for activation calibration
• PyO3 -- Rust/Python interop
• prost + protox -- pure-Rust ONNX protobuf pipeline (no system protoc required)