keypointrcnn_resnet50_fpn模型ONNX转TensorRT引擎失败求助

我在将keypointrcnn_resnet50_fpn模型从ONNX转换为TensorRT引擎时遇到诸多困难，经大量搜索仍无法生成引擎。

导出ONNX模型的代码 #

torch.onnx.export(model.cpu(),
                    input_tensor.cpu(),
                    onnx_file_path,
                    export_params = True,
                    do_constant_folding = False,
                    input_names = ['input'],
                    output_names = ['boxes', 'labels', 'scores', 'keypoints', 'keypoints_scores'],
                    dynamic_axes = {'input': {2 : 'height', 3 : 'width'}},
                    opset_version = 19
                    )

ONNX模型预处理及转换尝试步骤 #

# load the ONNX model
onnx_model = onnx.load(onnx_file_path)
# simplify the model
model_simp, check = simplify(onnx_model)
# export the simplified model
if check == True:
    onnx.save(model_simp, f"_simplified{onnx_file_name}")
else:
    print("ERROR: Failed to simplify and save model")

# re-export model suitable for TensorRT conversion
cmd = f"python3 -m onnxruntime.transformers.optimizer \
        --input=_simplified{onnx_file_name} \
        --output=_optimized{onnx_file_name} \
        "
subprocess.run(cmd, shell = True)
cmd = f"python3 -m onnxruntime.quantization.preprocess \
        --input=_optimized{onnx_file_name} \
        --output=_q_preprocessed{onnx_file_name} \
        "
subprocess.run(cmd, shell = True)

# re-export model with inferred shape
reloaded_model = onnx.load('_q_preprocessed_keypointrcnn_resnet50_fpn_o19.onnx')
onnx.checker.check_model(reloaded_model)
inferred_model = onnx.shape_inference.infer_shapes(reloaded_model, check_type = True, strict_mode = True, data_prop = True)
onnx.save(inferred_model, f"_shape_inferred{onnx_file_name}")

# ====

# check validity of model prior to conversion
cmd = f"polygraphy run _shape_inferred{onnx_file_name} --onnxrt"
subprocess.run(cmd, shell = True)

# ====

# generate TensorRT engine
cmd = f"~/.../TensorRT-10.0.0.6/bin/trtexec \
        --onnx=_shape_inferred{onnx_file_name} \
        --minShapes=input:1x3x512x512 \
        --optShapes=input:2x3x512x512 \
        --maxShapes=input:5x3x512x512 \
        --saveEngine=_final_keypointrcnn_resnet50_fpn.trt \
        --useCudaGraph \
        "
subprocess.run(cmd, shell = True)

转换TensorRT时的错误信息 #

[E] Error[4]: [shapeContext.cpp::operator()::3946] Error Code 4: Shape Error (reshape wildcard -1 has infinite number of solutions or no solution. Reshaping [0,8] to [0,-1,4].)
[E] [TRT] ModelImporter.cpp:826: While parsing node number 447 [Reshape -> "/roi_heads/Reshape_1_output_0"]:
[E] [TRT] ModelImporter.cpp:829: --- Begin node ---
    input: "/roi_heads/Flatten_output_0"
    input: "/roi_heads/Concat_2_output_0"
    output: "/roi_heads/Reshape_1_output_0"
    name: "/roi_heads/Reshape_1"
    op_type: "Reshape"
    attribute {
      name: "allowzero"
      i: 0
      type: INT
    }

问题定位 #

已用Netron查看模型架构，确认问题出在上述报错的Reshape节点区域：

解决建议 #

• 调整ONNX导出参数：
  1. 优先固定输入的高度和宽度，仅保留batch维度动态，目标检测模型的ROI分支对动态H/W的兼容性较差。修改导出代码示例：

  torch.onnx.export(model.cpu(),
                    input_tensor.cpu(),
                    onnx_file_path,
                    export_params=True,
                    do_constant_folding=True,  # 开启常量折叠减少动态节点
                    input_names=['input'],
                    output_names=['boxes', 'labels', 'scores', 'keypoints', 'keypoints_scores'],
                    dynamic_axes={'input': {0: 'batch'}},  # 仅开放batch维度动态
                    opset_version=17  # 降低opset版本提升TensorRT兼容性
                    )
  

• 简化预处理流程：
  移除onnxruntime.transformers.optimizer和量化预处理步骤，仅用onnx-simplify简化模型后直接尝试转换，多余的优化可能引入额外动态节点。
• 修复Reshape节点逻辑：
  若必须保留动态H/W，可手动修改ONNX模型，将Reshape节点的形状输入改为明确计算的结果（用Shape+Gather+Concat节点组合出目标形状），避免使用-1 wildcard导致TensorRT无法推导。
• 调整TensorRT转换参数：
  先移除--useCudaGraph参数确保基础转换成功，再逐步添加优化选项；确认输入的min/opt/max形状维度逻辑一致。

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