首先,无论是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之前进行插件注册,这样就可以使用自定义插件了。
#include "NvInferPlugin.h"
bool build() {
initLibNvInferPlugins(&gLogger, "");
...
}
Python中使用自定义插件(可以不下载源码):
- 还是刚刚的myplugin
- 在代码中添加REGISTER_TENSORRT_PLUGIN(MyPluginCreator)
- 自己写一个CMakeLists.txt,包括所有需要的文件和库,参考如下:
cmake_minimum_required(VERSION 3.8 FATAL_ERROR)
project(TensorRTPlugin LANGUAGES CXX CUDA)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wno-long-long -pedantic -Wno-deprecated-declarations")
set(PLUGIN_SOURCES)
set(PLUGIN_CU_SOURCES)
macro(set_ifndef var val)
if (NOT ${var})
set(${var} ${val})
endif()
message(STATUS "Configurable variable ${var} set to ${${var}}")
endmacro()
set(PLUGINS_LIST
MyPlugin
)
foreach(PLUGIN_ITER ${PLUGINS_LIST})
include_directories(${PLUGIN_ITER})
add_subdirectory(${PLUGIN_ITER})
endforeach(PLUGIN_ITER)
#源码中的plugin/common
add_subdirectory(common/)
set_ifndef(TRT_LIB /usr/lib/x86_64-linux-gnu/)
find_library(_NVINFER_LIB nvinfer HINTS ${TRT_LIB} PATH_SUFFIXES lib lib64)
set_ifndef(NVINFER_LIB ${_NVINFER_LIB})
include_directories(/usr/local/cuda-11.0/targets/x86_64-linux/include/)
list(APPEND PLUGIN_SOURCES "${PLUGIN_CU_SOURCES}")
add_library(tensorrtplugin MODULE ${PLUGIN_SOURCES})
target_link_libraries(tensorrtplugin PRIVATE ${NVINFER_LIB})
set_target_properties(tensorrtplugin PROPERTIES
CUDA_SEPARABLE_COMPILATION ON
)
- 编译生成libtensorrtplugin.so文件
- 在Python代码中添加库文件即可:
import ctypes
ctypes.CDLL("plugin/build/libtensorrtplugin.so")
...