Struct OnnxModel 

pub struct OnnxModel { /* private fields */ }

An ONNX model loaded from a protobuf file.

Provides methods for inspecting, extracting weights, saving quantized models, and validating graph connectivity.

Implementations

impl OnnxModel

pub fn load(path: impl AsRef<Path>) -> Result<Self>

Load an ONNX model from a file path.

Reads the entire file into a Vec<u8> before decoding. For multi-gigabyte models consider load_mmap (requires the mmap feature) to avoid the extra heap buffer.

Errors

Returns QuantizeError::ModelLoad if the file cannot be opened, is too large (>10 GB), or contains invalid protobuf data.

pub fn from_bytes(bytes: &[u8]) -> Result<Self>

Decode an ONNX model directly from a byte slice.

Useful for in-memory or fuzzing scenarios where the source isn’t a filesystem path. The same 10 GB size cap that load applies to files is also enforced here so callers feeding bytes from untrusted sources (HTTP, IPC, fuzz harnesses) can’t OOM the decoder with a pathologically large input.

Errors

Returns QuantizeError::ModelLoad if bytes exceeds 10 GB or cannot be decoded as a ModelProto.

pub fn info(&self) -> ModelInfo

Return a summary of the model’s structure.

pub fn input_shapes(&self) -> Vec<Vec<i64>>

Return the shapes of each graph input from the protobuf type info.

Each inner Vec<i64> contains the dimension values. Dynamic dims (symbolic or missing) are returned as -1. Returns one entry per graph.input that has tensor type information.

pub fn count_non_fp32_weight_initializers(&self) -> usize

Number of weight-shaped initializers whose dtype is a non-FP32 floating-point type (FP16, BF16, or Double). Useful for the CLI to explain why extract_weights returned nothing on a model that visibly has data — most commonly an FP16-exported HuggingFace model.

Only float-family dtypes are counted; rank-≥2 INT64 tensors are usually shape constants for Reshape / Tile / Gather and would otherwise show up as “non-FP32 weights” in the error message, confusing users.

pub fn count_external_data_initializers(&self) -> usize

Number of initializers whose tensor data lives in an external file (data_location == EXTERNAL), rather than inline in the protobuf.

quantize-rs reads only inline raw_data / float_data, so external-data tensors are skipped by extract_weights. The CLI and Python layers use this to turn an otherwise-confusing “no weight tensors found” into a precise diagnostic: ONNX exports above ~2 GB (large LLMs in particular) commonly store weights in a sidecar .onnx.data file, which must be inlined before quantization.

pub fn extract_weights(&self) -> Vec<WeightTensor>

Extract the quantizable FP32 weight tensors from the model’s initializers.

Only rank-≥2 tensors are returned. Rank-0/1 initializers — biases, BatchNorm scale/B/mean/var, LayerNorm parameters, PRelu slopes — are not weights and must not be quantized: per-tensor INT8 on a BatchNorm running_var rounds near-zero variances to 0, and the 1/sqrt(var) in BatchNorm then explodes the activations (this broke MobileNetV2 outright, cosine ≈ 0.10). Genuine quantization targets are always rank ≥ 2 (Conv 4-D, MatMul/Gemm 2-D, embedding tables 2-D).

QDQ scale scaffolding is also excluded: a {base}_scale FP32 initializer that has a sibling {base}_quantized is a DequantizeLinear scale (not a weight), as is any _quantize_rs_-prefixed initializer the save path synthesizes. Without this, loading an already-quantized model and quantizing it again would quantize the scales and silently corrupt the dequantization.

pub fn total_size_bytes(&self) -> usize

Total size of all weight tensors in bytes (float32).

Prefer computing this from already-extracted weights when available: weights.iter().map(|w| w.size_bytes()).sum() avoids reparsing.

impl OnnxModel

pub fn save_quantized( &mut self, quantized_data: &[QdqWeightInput], path: impl AsRef<Path>, ) -> Result<()>

Save a quantized model using the QDQ (DequantizeLinear) pattern.

Signature is identical to v0.2.0 — existing callers (CLI, calibration pipeline, examples) compile without changes.

What changed internally

v0.2.0 appended metadata to initializer names (e.g. conv1.weight → conv1.weight__qINT8_s0.001_z-3_len9408) without updating the nodes that reference them. ONNX Runtime rejected these models on load.

v0.3.0 inserts a DequantizeLinear node per weight. The node’s output carries the original name, so every downstream node is unchanged. Graph connectivity is preserved by construction, and the resulting model loads and runs in ONNX Runtime.

INT4 storage note

DequantizeLinear requires INT8 input in opsets < 21. By default, INT4-quantized values ([-8, 7]) are widened to INT8 bytes — 4× compression from FP32. For true 8× compression, call save_quantized_with_options with SaveOptions::with_native_int4(true), which emits native INT4 initializers and bumps the opset to 21.

pub fn save_quantized_with_options( &mut self, quantized_data: &[QdqWeightInput], path: impl AsRef<Path>, options: SaveOptions, ) -> Result<()>

Save a quantized model with explicit SaveOptions control.

See save_quantized for the transform details. Enabling SaveOptions::native_int4 for INT4 weights bumps the required opset to 21 automatically.

Fields not preserved on save

quantize-rs models a subset of the ONNX schema, so re-encoding drops any section outside it: ModelProto.functions (local-function custom ops), GraphProto.sparse_initializer, ModelProto.training_info, and assorted metadata_props/doc_string on nodes and tensors. When a loaded model carried functions, sparse_initializer, or training_info, this method prints a warning to stderr. Models built from such sections (notably custom-op graphs) should be quantized with care — the dequantized weights are correct, but the saved graph will not contain those sections.

impl OnnxModel

pub fn validate_connectivity(&self) -> ConnectivityReport

Check that every node input in the graph resolves to a known tensor.

A “known tensor” is one of:

• a declared graph input
• an initializer
• the output of a node appearing earlier in the node list

This is the exact check ONNX Runtime performs on load. It’s the check that v0.2.0’s validate command skipped, which is why the rename bug went undetected. Integrate report.summary() into the CLI validate output alongside the existing structure / weight checks.

impl OnnxModel

pub fn load_quantized_info(&self) -> Vec<QuantizedWeightInfo>

Extract metadata about quantized weights from a QDQ-format model.

Looks for initializer triples: {base}_quantized, {base}_scale, {base}_zp

Scale and zero-point are decoded in full — per-tensor yields a single element; per-channel yields one entry per channel. Bit-width comes from metadata_props (written by save_quantized); defaults to 8 if the metadata entry is missing.

Native INT4 zero-point tensors (DataType::Int4) are unpacked from their two-per-byte on-disk layout automatically.

Trait Implementations

impl Debug for OnnxModel

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.