Tensorflow 到 ONNX 中自定義操作的文檔

Question

我正在嘗試將經過訓練的 Tensorflow 2.5 SavedModel 轉換為 ONNX，希望最終能夠將 ONNX model 轉換為 TensorRT 以加速推理。 對於背景，它是我從 TensorFlow 1 到 TensorFlow 2 改編的3DFeatNet ( GitHub ) 論文的修改版本。

我的 model 需要在 Tensorflow 中使用自定義操作，它僅在 GPU 上運行。這些操作是使用tf.load_op_library()注冊的，它鏈接到編譯用於實現操作的 CUDA/C++ 代碼后生成的.so文件。

當我嘗試使用以下命令轉換 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. 有沒有辦法在 ONNX 中注冊.so文件，而無需將 CUDA/C++ 代碼調整為 ONNX 自定義操作 API？ 如果沒有，是否有關於如何操作的指南？
2. 當我將我的 model 轉換為 TensorRT 時，我是否需要經歷類似的過程來注冊自定義操作？

先感謝您。 如果需要更多信息，將編輯這篇文章。

Answer 1

通過將--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 特定於我正在使用的 model，並且可以根據您的用例中的給定模塊進行更改）

查看 Netron netron path/to/output/output_graph.onnx中生成的圖表顯示了我預期的結果，所以我現在關閉這個問題。

我仍然會感謝對我的方法的反饋，以及對我的第二個問題的回答。 謝謝！