Tensorrt自定义算子实现步骤

Tensorrt自定义算子实现步骤

目的

这篇文章主要是用于记录tensorrt自定义算子的过程.这里采用的是torch2trt直接转换的方式.

实验思路:

造一个pytorch中支持的操作,但是tensorrt中不支持的,也是说用torch2trt不能直接转换过去的算子,然后通过自定义一个插件实现跟torch中一样的功能,然后再用torch2trt接口就能转换成功了.

下面是具体步骤

在模型中造一个自定义的算子:

这里torch中的自定义算子为do_custom_op_shift,其功能是将输入参数增加一个偏移

定义一个模型,使用这个算子:

用minist数据集训练这个模型得到一个pytorch模型

用torch2trt接口转换模型,出现错误.

出现上面这种原因是因为在do_custom_op_shift这个函数中使用了tensorrt不支持的op,我们再看一下这个操作:

这个函数do_custom_op_shit在torch的模型中能够正常运行,但是不能转换成tensorrt的模型,现在需要做的就是给tensorrt增加一个自定义的算子,让他实现do_custom_op_shift的功能.

使用自定义算子解决问题

依赖工程

https://github.com/grimoire/torch2trt_dynamic
https://github.com/grimoire/amirstan_plugin

修改torch2trt_dynamic工程,这里是用python接口层面上增加一个自定义算子

再converters文件夹中增加一个converter命名为custom_op_shift.py,并在converters/init.py中增加一行

from .custom_op_shift import *

custom_op_shift.py的代码如下,这个文件的作用就是写一个接口,让他替换do_custom_op_shift的操作,这里的操作需要能够被tensorrt支持:

在plugins文件夹下增加一个plugin命名为create_shifop_plugin.py,并再plugins/init.py中增加一行代码

from .create_shiftop_plugin import create_shiftop_plugin

create_shifop_plugin.py的代码如下所示:

修改amirstan_plugin工程,这里是用ShiftPluginDynamic算子的具体实现

在include文件夹下增加一个文件,shiftPlugin.h

在src文件加下增加一个文件夹,命名为shiftPlugin,这里创建跟shiftPlugin有关的操作.

修改src/plugin/CMakeList.txt文件,增加shiftPlugin

修改src/plugin/InferPlugin.cpp,增加一行

重新编译这个工程.
重新安装torch2trt,再使用torch2trt接口转换上面的那个模型,就能成功了.