时间戳

首页 > 解决方案 > 从 TensorFlow 到 ONNX 的自定义操作文档

python - 从 TensorFlow 到 ONNX 的自定义操作文档

问题描述

我正在尝试将经过训练的 Tensorflow 2.5 SavedModel 转换为 ONNX，希望最终能够将 ONNX 模型转换为 TensorRT 以加速推理。作为背景，它是我从 TensorFlow 1 到 TensorFlow 2 改编的3DFeatNet ( GitHub ) 论文的修改版本。

我的模型需要在 Tensorflow 中使用自定义操作，这些操作仅在 GPU 上运行。这些操作使用 注册tf.load_op_library()，它链接到.so编译用于实现操作的 CUDA/C++ 代码后生成的文件。

当我尝试使用以下命令转换 SavedModel 时，自定义操作未注册，并且当我在 Netron 中查看文件时，生成的 ONNX 图表不显示自定义操作。

python tf2onnx_convert.py --saved-model path/to/model \ 
--output path/to/output \
--load_op_libraries path/to/.so_files \
--verbose --rename-inputs $INPUTS \
--rename-outputs $OUTPUTS

因此，我有两个问题：

  1. 有没有一种方法可以.so在 ONNX 中注册文件，而无需将 CUDA/C++ 代码调整为 ONNX 自定义操作 API？如果没有，是否有关于如何操作的指南？
  2. 当我将模型转换为 TensorRT 时，我是否需要在注册自定义操作时进行类似的过程？

先感谢您！如果需要更多信息，将编辑此帖子。

标签： pythontensorflowonnxtensorrt

解决方案

--custom-ops通过将标志添加到tf2onnx.convert命令中，我能够在稍微修改后回答我自己的问题。

我最终的命令变成了

python tf2onnx_convert.py --saved-model path/to/model \ 
--output path/to/output/output_graph.onnx \
--load_op_libraries path/to/.so_files \
--custom-ops QueryBallPoint,GroupPoint \
--rename-inputs $INPUTS \
--rename-outputs $OUTPUTS

（其中 QueryBallPoint 和 GroupPoint 特定于我正在使用的模型，并且可以根据您的用例中的给定模块进行更改）

查看 Netron 生成的图表netron path/to/output/output_graph.onnx显示了我的预期，所以我现在关闭这个问题。

我仍然会感谢对我的方法的反馈，以及对我的第二个问题的回答。谢谢！

推荐阅读