ONNX-Tensorrt如何注册自己的Op

ONNX-TensorRT自定义OP注册
在NVIDIA开源项目ONNX-TensorRT v5.0中,介绍如何为TX2的18.04系统Cuda10环境下注册自定义OP。以DCNV2为例,展示了编写源码、添加注册及在builtin_op_importers.cpp中定义importer的具体步骤。

前言

在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。
评论 7
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值