aitune.torch.backend.tensorrt.modelopt_calibration

NVIDIA ModelOpt Calibration module for quantization and autocast.

Module Contents

Functions

Name	Description
`_collect_per_sample_dicts`	Build list of per-sample dicts (ONNX name -> array), skip empty, raise on batch size 0.
`_concatenate_per_sample`	Merge list of per-sample dicts into one dict with arrays concatenated on axis=0.
`_filter_to_representative_max_shape`	Keep only samples whose non-batch shape matches the representative.
`_sample_shape_rank`	Return a rank for this sample (larger = larger shape). Used to pick max-shape representative.
`_sample_to_input_dict`	Convert one (args, kwargs) sample to a dict of ONNX input name -> numpy array.
`_shapes_match`	Return True if sample has the same non-batch shapes as reference for all keys.
`prepare_calibration_data`	Prepare calibration data in ModelOpt format from 1..N samples and graph_spec.

Data

logger

API

aitune.torch.backend.tensorrt.modelopt_calibration._collect_per_sample_dicts(
    data: list[aitune.torch.module.recording_module.Sample],
    graph_spec: aitune.torch.module.graph_spec.GraphSpec
) -> list[dict[str, numpy.ndarray]]

Build list of per-sample dicts (ONNX name -> array), skip empty, raise on batch size 0.

aitune.torch.backend.tensorrt.modelopt_calibration._concatenate_per_sample(
    per_sample: list[dict[str, numpy.ndarray]]
) -> dict[str, numpy.ndarray]

Merge list of per-sample dicts into one dict with arrays concatenated on axis=0.

aitune.torch.backend.tensorrt.modelopt_calibration._filter_to_representative_max_shape(
    per_sample: list[dict[str, numpy.ndarray]]
) -> list[dict[str, numpy.ndarray]]

Keep only samples whose non-batch shape matches the representative.

The representative is the sample with the largest total number of elements. Matching is on shape[1:] only, so different batch sizes are kept and concatenated along axis=0. Samples with different non-batch shapes are dropped with a warning.

aitune.torch.backend.tensorrt.modelopt_calibration._sample_shape_rank(
    sample: dict[str, numpy.ndarray]
) -> int

Return a rank for this sample (larger = larger shape). Used to pick max-shape representative.

aitune.torch.backend.tensorrt.modelopt_calibration._sample_to_input_dict(
    sample: aitune.torch.module.recording_module.Sample,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec
) -> dict[str, numpy.ndarray]

Convert one (args, kwargs) sample to a dict of ONNX input name -> numpy array.

aitune.torch.backend.tensorrt.modelopt_calibration._shapes_match(
    sample: dict[str, numpy.ndarray],
    reference: dict[str, numpy.ndarray]
) -> bool

Return True if sample has the same non-batch shapes as reference for all keys.

Compares shape[1:] only so different batch sizes (axis=0) are allowed and can be concatenated along axis=0.

aitune.torch.backend.tensorrt.modelopt_calibration.prepare_calibration_data(
    data: list[aitune.torch.module.recording_module.Sample],
    graph_spec: aitune.torch.module.graph_spec.GraphSpec
) -> dict[str, numpy.ndarray]

Prepare calibration data in ModelOpt format from 1..N samples and graph_spec.

Builds a single dict of arrays (one per ONNX input) by mapping each sample to ONNX input names and concatenating all samples along axis=0. So you get one “large batch” per input, e.g. 3 samples of shapes (2,3,224,224), (1,3,224,224), (4,3,224,224) → one array of shape (7, 3, 224, 224). ModelOpt requires this format (dict of arrays, not a list of batches).

Tensors are moved to CPU and converted to numpy. Uses graph_spec.input_spec for ONNX input names and locators into normalized forward arguments.

Parameters:

data

list[Sample]

List of 1..N Sample objects (args, kwargs); each can have any batch size.

graph_spec

GraphSpec

Graph specification whose input_spec defines ONNX input names and locators into normalized forward arguments.

Returns: dict[str, np.ndarray]

Single dict mapping ONNX input names to numpy arrays. Each value has

Raises:

ValueError: If data is empty or graph_spec has no tensor inputs.

aitune.torch.backend.tensorrt.modelopt_calibration.logger = logging.getLogger(__name__)