calib_utils

Provides basic calibration utils.

Classes

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

Functions

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)

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.

get_next()

Returns the next available calibration input from the reader.

class RandomDataProvider

Bases: CalibrationDataReader

Calibration data reader class with random data provider.

__init__(onnx_path)

Initializes the data reader class with random calibration data.

Parameters:

onnx_path (str) –

get_next()

Returns the next available calibration input from the reader.

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