想了解一下register_custom_op_symbolic自定义运算符的原理

dlunion / tensorRTIntegrate

TensorRT ONNX Plugin、Inference、Compile

464 stars 114 forks source link

Open 13213085 opened 4 years ago

dlunion commented 4 years ago

我没有使用过register_custom_op_symbolic这个，所以没有办法回答你这个问题。通常我是通过添加symbolic函数来实现自定义运算符。

13213085 commented 4 years ago

我最近在将mask rcnn从pytorch转到onnx，转完后发现，两个版本的输出不太一致，下面是nms的python接口，这里具体实现是用C++写的然后编译了进去，不知道能不能用symbolic函数实现

from maskrcnn_benchmark import _C import torch# we need this for the custom ops to exist import maskrcnn_benchmark._custom_ops # noqa: F401

from maskrcnn_benchmark.utils import amp

_nms = torch.ops.maskrcnn_benchmark.nms

nms = amp.float_function(_nms)

dlunion commented 4 years ago

symbolic+插件可以实现任何自定义需求，比如C++的。你可以参考我写的HSwish、DCN等

13213085 commented 4 years ago

好的，十分感谢你的帮助