首先,无论是python还是C++,我们都要写一个自己的plugin,一般的不支持op网上都有参考的版本,写plugin的参考教程:如何写好一个op
C++实现和使用过程:
- 下载TensorRT官网源码, 进入plugin可以看到许多已经定义好的插件可以参考,照着写就行,以ScatterND为例。
- 先写好自己的plugin,如名为:myplugin
注意:此处plugin代码里面不能包含REGISTER_TENSORRT_PLUGIN(MyPluginCreator)代码,不然会导致重复注册的问题。 - 在plugin/CMakeLists.txt添加自己大的plugin名称:
- 在plugin/InferPlugin.cpp添加注册信息
- 在路径下重新编译源代码
cd TensorRT
mkdir build
cd build
cmake .. -DTRT_LIB_DIR=$TRT_RELEASE/lib -DTRT_OUT_DIR=`pwd`/out
make -j${
nproc}
- 可以看到build路径下有个out文件夹,里面包括libnvinfer_plugin.so的三个文件。将这三个文件复制到自己的目录下,并删掉原来libnvinfer_plugin.so三个文件。
- 在自己的TensorRT推理代码的build model之前进行插件注册,这样就可以使用自定义插件了。
#