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, 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) –

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_model, calibration_shapes=None)

Initializes the data reader class with random calibration data.

Parameters:

  • onnx_model (str | ModelProto) –

  • calibration_shapes (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