TensorRt - caffe中支持prelu

好久没写博客了,今天就更新一下。

我们都知道TensorRT的激活层,支持relu,却不支持prelu。但是总有需要使用prelu层的时候。

补充: prelu层 相当于是带参数的relu

运行环境: ubuntu16.04+ TenosrRt 4.0.1以上 + cuda 8.0/cuda9.0 + Cudnn 7.1 + opencv3

一、下载

 先去github下载tensorrt-prelu的相关文件

并拷贝到我们的工程目录的源文件目录下

二、将Gplugin.cu重命名为GpluginGPU.cu

TensorRt - caffe中支持prelu_第1张图片

否则编译会出错

三、在GpluginGPU.h中 217行中DataType改为nvinfer1::DataType

 即

改为

四、添加prelu

TensorRt - caffe中支持prelu_第2张图片

在caffeToGIEModel函数结束的时候,要记得销毁parser

parser->destory();

调用caffeToGIEModel:

五、修改deploy文件

在deploy.protox中

将prelu层的type从prelu改为relu,prelu_param保留

TensorRt - caffe中支持prelu_第3张图片

 

这样就可以在tensorrt中使用prelu层了

 

 

你可能感兴趣的:(tensorrt)