onnx转TensorRT推理自定义插件(op)的方式,Python&C++

首先,无论是python还是C++,我们都要写一个自己的plugin,一般的不支持op网上都有参考的版本,写plugin的参考教程:如何写好一个op
C++实现和使用过程:

  1. 下载TensorRT官网源码, 进入plugin可以看到许多已经定义好的插件可以参考,照着写就行,以ScatterND为例。
  2. 先写好自己的plugin,如名为:myplugin
    注意:此处plugin代码里面不能包含REGISTER_TENSORRT_PLUGIN(MyPluginCreator)代码,不然会导致重复注册的问题。
  3. 在plugin/CMakeLists.txt添加自己大的plugin名称:onnx转TensorRT推理自定义插件(op)的方式,Python&C++_第1张图片
  4. 在plugin/InferPlugin.cpp添加注册信息
    onnx转TensorRT推理自定义插件(op)的方式,Python&C++_第2张图片
    onnx转TensorRT推理自定义插件(op)的方式,Python&C++_第3张图片
  5. 在路径下重新编译源代码
cd TensorRT
mkdir build
cd build
cmake  ..   -DTRT_LIB_DIR=$TRT_RELEASE/lib -DTRT_OUT_DIR=`pwd`/out
make -j${nproc}
  1. 可以看到build路径下有个out文件夹,里面包括libnvinfer_plugin.so的三个文件。将这三个文件复制到自己的目录下,并删掉原来libnvinfer_plugin.so三个文件。
  2. 在自己的TensorRT推理代码的build model之前进行插件注册,这样就可以使用自定义插件了。
#include "NvInferPlugin.h"
bool build() {	
    initLibNvInferPlugins(&gLogger, "");
    ...
}

Python中使用自定义插件(可以不下载源码):

  1. 还是刚刚的myplugin
  2. 在代码中添加REGISTER_TENSORRT_PLUGIN(MyPluginCreator)
  3. 自己写一个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
)
  1. 编译生成libtensorrtplugin.so文件
  2. 在Python代码中添加库文件即可:
import ctypes
ctypes.CDLL("plugin/build/libtensorrtplugin.so")
...

你可能感兴趣的:(TensorRT,onnx,C++,api,python,c++,深度学习)