trt_utils

This module contains TensorRT utils.

Functions

`get_custom_layers`	Gets custom layers in ONNX model.
`infer_types_shapes`	Updates tensor shapes in ORT graph.
`load_onnx_model`	Load ONNX model.

get_custom_layers(onnx_path, trt_plugins)

Gets custom layers in ONNX model.

Parameters:

onnx_path (str) – Path to the input ONNX model.
trt_plugins (str | None) – Paths to custom TensorRT plugins.

Returns:

List of custom layers. Dictionary containing tensors information: {‘tensor_name’: {‘shape’: tensor.shape, ‘dtype’: tensor.dtype}}

Return type:

Tuple[List[str], Dict]

infer_types_shapes(graph, all_tensor_info)

Updates tensor shapes in ORT graph.

Parameters:

graph (Graph) – ONNX model’s GS graph.
all_tensor_info (Dict) – Dictionary containing tensors information.

Returns:

None. In-memory modification of graph.

Return type:

None

load_onnx_model(onnx_path, trt_plugins=None, override_shapes=None, use_external_data_format=False, intermediate_generated_files=None)

Load ONNX model. If ‘tensorrt’ is installed, check if the model has custom ops and ensure it’s supported by ORT.

Parameters:

onnx_path (str) – Path to the input ONNX model.
trt_plugins (str | None) – Paths to custom TensorRT plugins.
override_shapes (str) – Override model input shapes with static shapes.
use_external_data_format (bool) – If True, separate data path will be used to store the weights of the quantized model.
intermediate_generated_files (list[str]) – List of paths of intermediate ONNX files, generated during quantization.

Returns:

Loaded ONNX model supported by ORT. Boolean indicating whether the model has custom ops or not. List of custom ops in the ONNX model. Path to new intermediary ONNX model.

Return type:

tuple[ModelProto, bool, list[str], str]