calib_utils

Provides calibration data readers and utilities.

Classes

CalibrationDataProvider	Calibration data provider class.
NpyCalibrationReader	Stream single-input calibration batches from a directory of NPY files.
NpzCalibrationReader	Stream single- or multi-input calibration batches from a directory of NPZ files.
RandomDataProvider	Calibration data reader class with random data provider.

Functions

create_directory_calibration_reader	Create a bounded calibration reader for a directory of NPY or NPZ batches.
import_scales_from_calib_cache	Reads TensorRT calibration cache and returns as dictionary.

class CalibrationDataProvider

Bases: CalibrationDataReader

Calibration data provider class.

__init__(onnx_path, calibration_data, calibration_shapes=None)

Initializes the data provider class with the calibration data iterator.

Parameters:

• onnx_path (str | ModelProto) – Path to the ONNX model.

• calibration_data (ndarray | dict[str, ndarray]) – Numpy data to calibrate the model. Ex. If a model has input shapes like {“sample”: (2, 4, 64, 64), “timestep”: (1,), “encoder_hidden_states”: (2, 16, 768)}, the calibration data should have dictionary of tensors with shapes like {“sample”: (1024, 4, 64, 64), “timestep”: (512,), “encoder_hidden_states”: (1024, 16, 768)} to calibrate with 512 samples.

• calibration_shapes (str | None) – A string representing the shape of each input tensors for one calibration step. If the shape is not provided for an input tensor, the shape is inferred from the onnx model directly, with all the unknown dims filled with 1.

get_first()

Returns the first calibration input from the reader without incrementing the iterator.

This is useful when doing a test run for the session.

get_next()

Returns the next available calibration input from the reader.

rewind()

Rewinds the data reader to the first index.

class NpyCalibrationReader

Bases: _FileCalibrationReader

Stream single-input calibration batches from a directory of NPY files.

__init__(calibration_dir, onnx_path, input_name=None, max_batches=512, max_batch_bytes=134217728, max_tensor_elements=32000000, max_onnx_bytes=536870912, calibration_shapes=None)

Initialize an NPY reader from the model input metadata.

load(batch_path)

Validate and load one NPY calibration batch.

class NpzCalibrationReader

Bases: _FileCalibrationReader

Stream single- or multi-input calibration batches from a directory of NPZ files.

__init__(calibration_dir, onnx_path, max_batches=512, max_archive_bytes=134217728, max_tensor_elements=32000000, max_onnx_bytes=536870912, safe_cast_inputs=(), calibration_shapes=None)

Initialize an NPZ reader from the model input metadata.

load(batch_path)

Validate and load one NPZ calibration batch.

class RandomDataProvider

Bases: CalibrationDataReader

Calibration data reader class with random data provider.

__init__(onnx_model, calibration_shapes=None)

Initializes the data reader class with random calibration data.

Parameters:

• onnx_model (str | ModelProto)

• calibration_shapes (str | None)

get_first()

Returns the first calibration input from the reader without incrementing the iterator.

This is useful when doing a test run for the session.

get_next()

Returns the next available calibration input from the reader.

rewind()

Rewinds the data reader to the first index.

create_directory_calibration_reader(calibration_dir, onnx_path, *, max_batches=512, calibration_shapes=None, safe_cast_inputs=())

Create a bounded calibration reader for a directory of NPY or NPZ batches.

import_scales_from_calib_cache(cache_path)

Reads TensorRT calibration cache and returns as dictionary.

Parameters:

cache_path (str) – Calibration cache path.

Returns:

float_scale}.

Return type:

Dictionary with scales in the format {tensor_name