CODEWITHSUNDEEP
pythontensorflowonnxtensorrt
tianyilim
tianyilim

Reputation: 21

Documentation for Custom Ops in Tensorflow to ONNX

I am attempting to convert a trained Tensorflow 2.5 SavedModel into ONNX, with the hopes of eventually being able to convert the ONNX model into TensorRT to accelerate inference. For background, it is a modified version of the 3DFeatNet (GitHub) paper that I adapted from TensorFlow 1 to TensorFlow 2.

My model requires the use of custom ops in Tensorflow, which run only on GPU. These ops are registered using tf.load_op_library(), which links to the .so file generated after compiling the CUDA/C++ code used to implement the op.

When I attempt to convert the SavedModel using the below command, the custom ops are not registered, and the resulting ONNX graph does not show the custom ops when I view the file in Netron.

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

Thus, I have 2 questions:

  1. Is there a way to register the .so files in ONNX without needing to adapt the CUDA/C++ code to the ONNX custom op API? If not, are there any guides on how to do so?
  2. Will I need to undergo a similar process in registering custom ops when I convert my model to TensorRT?

Thank you in advance! Will edit this post if any more info is needed.

Upvotes: 0

Views: 880

Answers (1)

tianyilim
tianyilim

Reputation: 21

I was able to answer my own question after tinkering around a bit by adding the --custom-ops flag to the tf2onnx.convert command.

My eventual command became

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

(where QueryBallPoint and GroupPoint are specific to the model I am using, and could be changed as per the given module that is in your use case)

Viewing the graph produced in Netron netron path/to/output/output_graph.onnx shows something I expected, so I'm closing the issue now.

I would still appreciate feedback on my method, and an answer to my second question. Thanks!

Upvotes: 2

Related Questions