trt_utils

此模块包含TensorRT工具。

函数

get_custom_layers

获取ONNX模型中的自定义层。

load_onnx_model

加载ONNX模型。
get_custom_layers(onnx_path, trt_plugins)

获取ONNX模型中的自定义层。

Parameters:

  • onnx_path (str) – 输入ONNX模型的路径。

  • trt_plugins (str | None) – 自定义TensorRT插件的路径。

Returns:

自定义层列表。

Return type:

列表[字符串]

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

加载ONNX模型。如果安装了‘tensorrt’，检查模型是否有自定义操作，并确保它被ORT支持。

Parameters:

  • onnx_path (str) – 输入ONNX模型的路径。

  • trt_plugins (str | None) – 自定义TensorRT插件的路径。

  • use_external_data_format (bool) – 如果为True，将使用单独的数据路径来存储量化模型的权重。

Returns:

加载由ORT支持的ONNX模型。 布尔值，指示模型是否具有自定义操作。 ONNX模型中的自定义操作列表。

Return type:

元组[模型原型, 布尔值, 列表[字符串]]