onnx转tensorRT模型出现错误 This version of TensorRT only supports input K as an initializer

问题

onnx模型转tensorRT模型时，出现错误。

This version of TensorRT only supports input K as an initializer. Try applying constant folding on the model using Polygraph

google到tensorRT 8.6支持了dynamic topk，不会再有这个问题。
但项目上限制是 tensorRT 8.5 Problems converting keypoint RCNN from Detectron2 to TensorRT · Issue #2678 · NVIDIA/TensorRT

对比出错处的topk算子，可以看到正常转tensorRT的topK算子是没有Identity输入的。

解决方案

借助 onnx_graphsurgeon库将topK算子的Identity输入强制转换为topK的Constant。
脚本

import onnx
import onnx_graphsurgeon as gs
import numpy as np

# 加载 ONNX 模型
model_path = 'model.onnx'
onnx_model = onnx.load(model_path)

# 使用 Polygraphy 进行常量折叠
folded_model = fold_constants(onnx_model)

# 使用 onnx_graphsurgeon 将 TopK 的 k 转换为常量
graph = gs.import_onnx(onnx_model)
#graph = gs.import_onnx(folded_model)
for node in graph.nodes:
    if node.op == 'TopK' :
        print(node)
        k_input = node.inputs[1]
        if k_input.inputs and isinstance(k_input.inputs[0], gs.ir.node.Node):
            identity_node = k_input.inputs[0]
            node.inputs[1] = identity_node.inputs[0]

# 导出修改后的模型
modified_model_path = 'output.fold.onnx'
onnx.save(gs.export_onnx(graph), modified_model_path)
#onnx.save(folded_model, modified_model_path)

# 检查模型
onnx.checker.check_model(modified_model_path)
print("Model checked successfully!")