ONNX-Tensorrt如何注册自己的Op

前言

在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中。

注册具体步骤

  1. 编写源码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);
  1. 最核心的步骤,builtin_op_importers.cpp中使用DEFINE_BUILTIN_OP_IMPORTER添加对自己注册Op的使用,只有这里生效,使用onnx2trt这个程序的时候,才能找到我们自己定制的Op。
发布了45 篇原创文章 · 获赞 21 · 访问量 5万+

猜你喜欢

转载自blog.csdn.net/xxradon/article/details/94858139