导出支持动态批次的yolov8m.engine遇TensorRT Range算子错误求替代方案

尝试将预训练YOLOv8m模型导出为支持动态批次的TensorRT .engine 文件，操作步骤及报错如下：

1. 按Ultralytics官方指令导出ONNX：

from ultralytics import YOLO

# 加载模型
model = YOLO('yolov8m.pt')  # 官方预训练模型
# model = YOLO('path/to/best.pt')  # 自定义训练模型

# 导出动态批次ONNX
model.export(format='onnx', dynamic=True)

2. 执行trtexec转换命令：

trtexec --onnx=yolov8m.onnx --workspace=8144 --fp16 --minShapes=input:1x3x640x640 --optShapes=input:2x3x640x640 --maxShapes=input:10x3x640x640 --saveEngine=my.engine

3. 出现错误：

[08/10/2023-23:53:10] [I] TensorRT version: 8.2.5
[08/10/2023-23:53:11] [I] [TRT] [MemUsageChange] Init CUDA: CPU +336, GPU +0, now: CPU 348, GPU 4361 (MiB)
[08/10/2023-23:53:11] [I] [TRT] [MemUsageSnapshot] Begin constructing builder kernel library: CPU 348 MiB, GPU 4361 MiB
[08/10/2023-23:53:12] [I] [TRT] [MemUsageSnapshot] End constructing builder kernel library: CPU 483 MiB, GPU 4393 MiB
[08/10/2023-23:53:12] [I] Start parsing network model
[08/10/2023-23:53:12] [I] [TRT] ----------------------------------------------------------------
[08/10/2023-23:53:12] [I] [TRT] Input filename:   yolov8m.onnx
[08/10/2023-23:53:12] [I] [TRT] ONNX IR version:  0.0.8
[08/10/2023-23:53:12] [I] [TRT] Opset version:    17
[08/10/2023-23:53:12] [I] [TRT] Producer name:    pytorch
[08/10/2023-23:53:12] [I] [TRT] Producer version: 2.0.1
[08/10/2023-23:53:12] [I] [TRT] Domain:           
[08/10/2023-23:53:12] [I] [TRT] Model version:    0
[08/10/2023-23:53:12] [I] [TRT] Doc string:       
[08/10/2023-23:53:12] [I] [TRT] ----------------------------------------------------------------
[08/10/2023-23:53:12] [W] [TRT] onnx2trt_utils.cpp:366: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
[08/10/2023-23:53:12] [E] [TRT] ModelImporter.cpp:773: While parsing node number 305 [Range -> "/model.22/Range_output_0"]:
[08/10/2023-23:53:12] [E] [TRT] ModelImporter.cpp:774: --- Begin node ---
[08/10/2023-23:53:12] [E] [TRT] ModelImporter.cpp:775: input: "/model.22/Constant_8_output_0"
input: "/model.22/Cast_output_0"
input: "/model.22/Constant_9_output_0"
output: "/model.22/Range_output_0"
name: "/model.22/Range"
op_type: "Range"

[08/10/2023-23:53:12] [E] [TRT] ModelImporter.cpp:776: --- End node ---
[08/10/2023-23:53:12] [E] [TRT] ModelImporter.cpp:779: ERROR: builtin_op_importers.cpp:3353 In function importRange:
[8] Assertion failed: inputs.at(0).isInt32() && "For range operator with dynamic inputs, this version of TensorRT only supports INT32!"
[08/10/2023-23:53:12] [E] Failed to parse onnx file
[08/10/2023-23:53:12] [I] Finish parsing network model
[08/10/2023-23:53:12] [E] Parsing model failed
[08/10/2023-23:53:12] [E] Failed to create engine from model.

已知升级TensorRT版本可解决，但需寻找替代方案。

方案1：修改ONNX导出参数，避免类型异常 #

调整导出代码，关闭ONNX自动简化并指定兼容的opset版本，强制使用INT32类型：

from ultralytics import YOLO

model = YOLO('yolov8m.pt')
model.export(
    format='onnx',
    dynamic=True,
    simplify=False,  # 关闭自动简化，防止Range节点类型被修改
    opset=16,  # 使用兼容TensorRT 8.2.5的opset版本
    int8=False  # 保持FP32/FP16类型，避免额外类型转换问题
)

方案2：手动修改ONNX模型的Range节点输入类型 #

使用onnx库直接修改模型中Range节点的输入数据类型为INT32：

import onnx

# 加载原ONNX模型
model = onnx.load('yolov8m.onnx')

# 遍历所有节点，修复Range节点的INT64输入
for node in model.graph.node:
    if node.op_type == 'Range':
        # 遍历Range节点的前两个输入（start和limit）
        for input_idx in [0, 1]:
            input_name = node.input[input_idx]
            # 找到对应的初始值Constant节点
            for init in model.graph.initializer:
                if init.name == input_name and init.data_type == onnx.TensorProto.INT64:
                    # 转换数据类型为INT32
                    init.data_type = onnx.TensorProto.INT32
                    # 转换数据内容
                    init.int32_data[:] = [int(val) for val in init.int64_data]
                    # 清空原INT64数据
                    del init.int64_data[:]

# 保存修复后的模型
onnx.save(model, 'yolov8m_fixed.onnx')

修复后使用yolov8m_fixed.onnx重新执行原trtexec命令。

方案3：用ONNX Runtime工具自动修复兼容性 #

使用ONNX Runtime内置工具自动转换INT64节点为INT32：

python -m onnxruntime.tools.convert_onnx_models_to_trt_compatible yolov8m.onnx

工具会生成兼容TensorRT的ONNX文件，之后直接用该文件执行trtexec转换即可。

内容的提问来源于stack exchange，提问作者Pe Dro