贤者之路， Caffe转TensorRT

 

【引】将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