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接口转换上面的那个模型,就能成功了.