trt_utils

This module contains TensorRT utils.

Functions

get_custom_layers

Gets custom layers in ONNX model.

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.

Return type:

List[str]

load_onnx_model(onnx_path, trt_plugins=None, use_external_data_format=False)

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.

  • use_external_data_format (bool) – If True, separate data path will be used to store the weights of the quantized model.

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.

Return type:

Tuple[ModelProto, bool, List[str]]