前言
在NVIDIA开源项目中,onnx-tensorrt,版本是v5.0,目前TX2的18.04系统Cuda10,Tensorrt5.0.26目前只支持到这个版本。
我们可以看到tensorrt是如何将ONNX的模型转化为tensorrt支持的序列化模型。源码中我们可以看到有4个Op:Split,FancyActivation,InstanceNormalization,ResizeNearest都是后来注册支持的,并不是tensorrt内部支持的。
所以我们自己注册自己的custom Op,就需要参考这些op如何实现注册功能的。
注册自己的custom Op:dcn_v2_forward
DCN V2是可变形卷积的最新版本实现,DCN有很多优秀的特性,特别是目标检测算法中,加入DCN可以有效提高检出率。目前DCNV2版本的源码https://github.com/CharlesShang/DCNv2,目前具有Cuda加速实现,所以我们以DCN V2作为一个定制的Op注册到ONNX-Tensorrt中。
注册具体步骤
- 编写源码dcn_v2_forward.hpp和dcn_v2_forward.cu,这里面就是dcn_v2_forward的Op具体实现,重载onnx2trt::Plugin,注意这里是不同于tensorrt的caffe注册layer或者是tensorrt自己注册层。
class dcn_v2_forwardPlugin final : public onnx2trt::Plugin
具体编写步骤可以参考Split,FancyActivation,InstanceNormalization,ResizeNearest的实现。
3. 在builtin_plugins.cpp最后面,像下面一样添加对dcn_v2_forward的注册。
REGISTER_BUILTIN_PLUGIN("FancyActivation", FancyActivationPlugin);
REGISTER_BUILTIN_PLUGIN("ResizeNearest", ResizeNearestPlugin);
REGISTER_BUILTIN_PLUGIN("Split" , SplitPlugin);
REGISTER_BUILTIN_PLUGIN("InstanceNormalization", InstanceNormalizationPlugin);
REGISTER_BUILTIN_PLUGIN("dcn_v2_forward", dcn_v2_forwardPlugin);
REGISTER_BUILTIN_NVPLUGIN("Concat", ConcatPlugin);
- 最核心的步骤,builtin_op_importers.cpp中使用DEFINE_BUILTIN_OP_IMPORTER添加对自己注册Op的使用,只有这里生效,使用onnx2trt这个程序的时候,才能找到我们自己定制的Op。