calib_utils

Provides basic calibration utils.

Classes

`CalibrationDataProvider`	Calibration data provider class.
`RandomDataProvider`	Calibration data reader class with random data provider.

Functions

Reads TensorRT calibration cache and returns as dictionary.

class CalibrationDataProvider

Bases: CalibrationDataReader

Calibration data provider class.

__init__(onnx_path, calibration_data, calibration_shapes=None)

Intializes the data provider class with the calibration data iterator.

Parameters:

onnx_path (str) – 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) – 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.

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:

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.

import_scales_from_calib_cache(cache_path)

Reads TensorRT calibration cache and returns as dictionary.