【引】将Caffe转TensorRT的时候,有很多自己设计的接口TensorRT库本身不支持。我们需要自己创建Plugin,本文介绍TensorRT的创建,如何自定义Plugin,和快速书写cuda函数。
【结构】
将Caffe转TensorRT的时候,有很多自己设计的接口TensorRT库本身不支持。我们需要继承TensorRT里面的IPlugin类来创建自己的Plugin。然后ICaffeParser对象会通过IPluginFacotry 解析Prototext里的网络结构,参数等信息,并将解析的信息放在内存中以便之后的计算。同时IBuilder负责设置网络数据类型(int8 FP16)batchSzie等一些通用信息。这样一个网络就被建立,网络的结构,参数都被创建在内存之中被INetworkDefinition对象所管理
【重载plugin】
class UpsampleLayer : public IPlugin
{
public:
UpsampleLayer(){}
//serilize constructure
UpsampleLayer(size_t stride):stride_(stride)
{
}
//deserilize constructure
UpsampleLayer(const void* buffer,size_t sz, size_t stride):stride_(stride)
{
/* 反序列化, 将连续放置的内存,读取到成员变量里面 */
const int* d = reinterpret_cast<const int*>(buffer);
channel_=d[0];
w_=d[1];
h_=d[2];
}
// get output size infomation
inline int getNbOutputs() const override { return 1; };
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override
{
channel_=inputs[0].d[0];
w_=inputs[0].d[1];
h_=inputs[0].d[2];
return DimsCHW(inputs[0].d[0], inputs[0].d[1]*stride_, inputs[0].d[2]*stride_);
}
i