torch_tensorrt.ptq

class torch_tensorrt.ptq.CacheCalibrator(*args: Any, **kwargs: Any)

Bases: object

Constructs a calibrator class in TensorRT which directly uses pre-existing cache file for calibration. :param cache_file: path to cache file. :param algo_type: choice of calibration algorithm.

class torch_tensorrt.ptq.CalibrationAlgo(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: Enum

ENTROPY_CALIBRATION = <CalibrationAlgo.ENTROPY_CALIBRATION: 1>

ENTROPY_CALIBRATION_2 = <CalibrationAlgo.ENTROPY_CALIBRATION_2: 2>

LEGACY_CALIBRATION = <CalibrationAlgo.LEGACY_CALIBRATION: 0>

MINMAX_CALIBRATION = <CalibrationAlgo.MINMAX_CALIBRATION: 3>

class torch_tensorrt.ptq.DataLoaderCalibrator(*args: Any, **kwargs: Any)

Bases: object

Constructs a calibrator class in TensorRT and uses pytorch dataloader to load/preproces data which is passed during calibration. :param dataloader: an instance of pytorch dataloader which iterates through a given dataset. :param algo_type: choice of calibration algorithm. :param cache_file: path to cache file. :param use_cache: flag which enables usage of pre-existing cache. :param device: device on which calibration data is copied to.

torch_tensorrt.ptq.get_batch(self: object, _: Any) → Optional[List[int]]

torch_tensorrt.ptq.get_batch_size(self: object) → int

torch_tensorrt.ptq.get_cache_mode_batch(self: object) → None

torch_tensorrt.ptq.read_calibration_cache(self: object) → bytes

torch_tensorrt.ptq.write_calibration_cache(self: object, cache: bytes) → None

Classes

class torch_tensorrt.ptq.DataLoaderCalibrator(*args: Any, **kwargs: Any)

Constructs a calibrator class in TensorRT and uses pytorch dataloader to load/preproces data which is passed during calibration. :param dataloader: an instance of pytorch dataloader which iterates through a given dataset. :param algo_type: choice of calibration algorithm. :param cache_file: path to cache file. :param use_cache: flag which enables usage of pre-existing cache. :param device: device on which calibration data is copied to.

__init__(**kwargs: Any)

class torch_tensorrt.ptq.CacheCalibrator(*args: Any, **kwargs: Any)

Constructs a calibrator class in TensorRT which directly uses pre-existing cache file for calibration. :param cache_file: path to cache file. :param algo_type: choice of calibration algorithm.

__init__(**kwargs: Any)

Enums

class torch_tensorrt.ptq.CalibrationAlgo(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)