实现TensorRT-7.0插件自由!(如果不踩坑使用TensorRT插件功能) (2)

以MyCustomPlugin为例,假设我们的这个MyCustomPlugin有两个权重weight和bias以及一个参数in_channel(这个权重和参数没有啥意义,纯粹,纯粹为了演示):

private: int _in_channel; // 参数 std::vector<float> _weight; // 权重,在cpu空间存放 std::vector<float> _bias; // 偏置权重,在cpu空间存放 float* _d_weight; // 权重,在GPU空间存放 float* _d_bias; bool _initialized; cudnnHandle_t _cudnn_handle; const char* mPluginNamespace; std::string mNamespace; 构造函数和析构函数

构造函数一般设置为三个。

第一个用于在parse阶段,PluginCreator用于创建该插件时调用的构造函数,需要传递权重信息以及参数。

第二个用于在clone阶段,复制这个plugin时会用到的构造函数。

第三个用于在deserialize阶段,用于将序列化好的权重和参数传入该plugin并创建爱你哦。

以我们的MyCustomPlugin为例:

MyCustomPlugin(int in_channel, nvinfer1::Weights const& weight, nvinfer1::Weights const& bias); MyCustomPlugin(float in_channel, const std::vector<float>& weight, const std::vector<float>& bias); MyCustomPlugin(void const* serialData, size_t serialLength);

析构函数则需要执行terminate,terminate函数就是释放这个op之前开辟的一些显存空间:

MyCustomPlugin::~MyCustomPlugin() { terminate(); }

注意需要把默认构造函数删掉:

MyCustomPlugin() = delete; getNbOutputs

插件op返回多少个Tensor,比如MyCustomPlugin这个操作只输出一个Tensor(也就是一个output),所以直接return 1:

// MyCustomPlugin returns one output. int MyCustomPlugin::getNbOutputs() const { return 1; } initialize

初始化函数,在这个插件准备开始run之前执行。

主要初始化一些提前开辟空间的参数,一般是一些cuda操作需要的参数(例如conv操作需要执行卷积操作,我们就需要提前开辟weight和bias的显存),假如我们的算子需要这些参数,则在这里需要提前开辟显存。

需要注意的是,如果插件算子需要开辟比较大的显存空间,不建议自己去申请显存空间,可以使用Tensorrt官方接口传过来的workspace指针来获取显存空间。因为如果这个插件被一个网络调用了很多次,而这个插件op需要开辟很多显存空间,那么TensorRT在构建network的时候会根据这个插件被调用的次数开辟很多显存,很容易导致显存溢出。

getOutputDataType

返回结果的类型,一般来说我们插件op返回结果类型与输入类型一致:

nvinfer1::DataType InstanceNormalizationPlugin::getOutputDataType( int index, const nvinfer1::DataType* inputTypes, int nbInputs) const { ASSERT(inputTypes && nbInputs > 0 && index == 0); return inputTypes[0]; } getWorkspaceSize

这个函数需要返回这个插件op需要中间显存变量的实际数据大小(bytesize),这个是通过TensorRT的接口去获取,是比较规范的方式。

我们需要在这里确定这个op需要多大的显存空间去运行,在实际运行的时候就可以直接使用TensorRT开辟好的空间而不是自己去申请显存空间。

size_t MyCustomPlugin::getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const { // 计算这个op前向过程中你认为需要的中间显存数量 size_t need_num; return need_num * sizeof(float); } enqueue

内容版权声明:除非注明,否则皆为本站原创文章。

转载注明出处:https://www.heiqu.com/wpwjpf.html