TensorFlow自定义算子转ONNX及后续TensorRT转换问题咨询

问题1解答 #

不存在无需适配ONNX自定义算子API、直接复用TensorFlow侧.so文件注册ONNX算子的方法。
ONNX Runtime和TensorFlow的算子注册体系、调用接口、上下文管理逻辑完全独立不兼容，TF侧编译的.so只能适配TF的算子调用规范，无法被ONNX Runtime识别加载。
操作参考步骤如下：

• 首先补充tf2onnx的自定义算子映射逻辑：你之前使用的--load_op_libraries参数仅能让tf2onnx识别TF侧的算子结构避免转换报错，不会自动完成算子到ONNX的映射。你需要额外编写tf2onnx自定义转换脚本，把你的TF自定义算子的属性、输入输出参数一一映射为ONNX自定义算子节点，算子名、自定义域名要和后续注册的ONNX算子完全对齐，做完这一步转换出来的ONNX文件里就能看到对应的自定义算子节点了。
• 之后将现有CUDA/C++的Kernel代码适配ONNX Runtime的Custom Op API，编译为ONNX可用的.so文件，注册时匹配之前定义的算子名和域名，就可以在ONNX Runtime中加载调用。

问题2解答 #

将ONNX转换为TensorRT模型时，同样需要执行自定义算子的注册流程，且无法直接复用TF或ONNX侧的.so文件，需要适配TensorRT的Plugin接口。
针对你最终要落地TensorRT推理的场景，有两种适配路径可选：

• 路径1：完成ONNX侧的自定义算子适配后，再把CUDA Kernel代码适配TensorRT的IPluginV2系列接口，编写对应插件注册到TensorRT的ONNX解析器中，转模型时解析器会自动识别ONNX里的自定义算子调用对应插件。
• 路径2：跳过ONNX Runtime侧的算子适配步骤，在tf2onnx映射算子时直接给节点打上符合TensorRT解析规则的自定义标记，只需要编写一次TensorRT Plugin注册到解析器，就可以完成转换和推理，减少重复适配工作量。

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