1 This version of TensorRT only supports input K as an initializer

今天使用TensorRT-8.5.1.7的trtexec工具转换onnx模型为trt模型时出现了以下错误:

[04/27/2023-09:04:12] [E] [TRT] ModelImporter.cpp:729: --- End node ---
[04/27/2023-09:04:12] [E] [TRT] ModelImporter.cpp:732: ERROR: ModelImporter.cpp:168 In function parseGraph:
[6] Invalid Node - TopK_573
This version of TensorRT only supports input K as an initializer. Try applying constant folding on the model using Polygraphy: https://github.com/NVIDIA/TensorRT/tree/master/tools/Polygraphy/examples/cli/surgeon/02_folding_constants
[04/27/2023-09:04:12] [E] Failed to parse onnx file

从错误信息上看,是TopK算子没有找到,然后TensorRT给的建议是使用它的Polygraphy工具对Onnx模型进行constant folding。

之后我按照https://github.com/NVIDIA/TensorRT/tree/master/tools/Polygraphy/examples/cli/surgeon/02_folding_constants配置了Polygraphy工具的环境,并对onnx模型进行了constant folding操作,再次进行转换还是出现了问题。

2 问题解决方法

目前TensorRT版本已经到了8.6.0的预览版,我从8.6.0版本的官方文档https://docs.nvidia.com/deeplearning/tensorrt/archives/index.html#trt_8 看到了其的算子支持是包含TopK的,然后我使用了最新的8.6.0预览版本8.6.0.12重新对原始的onnx模型进行转换,竟然成功了。

TensorRT – Invalid Node – TopK,This version of TensorRT only supports input K as an initializer-StubbornHuang Blog

然后我看8.5.3的文档,里面竟然也是支持TopK算子的,然后我又重新使用了8.5.3对原始的onnx模型进行了转换,但是在解析模型时同样出现了TopK算子有错的信息。

这真是令人费解。

不过以后遇到TensorRT转换Onnx模型出现算子不支持的情况,首先应该查看最新TensorRT的算子支持列表,查看是否有这个算子,避免浪费太多时间。