TensorRT--为caffeparser添加自定义层支持

对于我们自定义的层如果写到了caffe prototxt中，在部署模型时调用caffeparser来解析就会报错。还是以Upsample为例，如果在prototxt中有下面这段来添加了一个upsample的层

layer {
  name: "upsample0"
  type: "Upsample"
  bottom: "ReLU11"
  top: "Upsample1"
}

这时再调用

const IBlobNameToTensor *blobNameToTensor =	parser->parse(deployFile.c_str(),
                                                              modelFile.c_str(),
                                                              *network,
                                                              modelDataType);

就会出现错误

could not parse layer type Upsample

之前我们已经编写了Upsample的插件，怎么让tensorRT的caffe parser识别出prototxt中的upsample层自动构建我们自己编写的插件呢？这时我们就需要定义一个插件工程类继承基类nvinfer1::IPluginFactory, nvcaffeparser1::IPluginFactoryExt。

class PluginFactory : public nvinfer1::IPluginFactory, public nvcaffeparser1::IPluginFactoryExt
{
    //其中必须要的实现的方法有:
    //判断一个层是否是plugin的方法，
    //输入的参数就是prototxt中layer的name，通过name来判断一个层是否注册为插件
    bool isPlugin(const char *name) override {
        return isPluginExt(name);
    }
    //判断层名字是否是upsample层的名字
    bool isPluginExt(const char *name) override {
        
        char *aa = new char[6];
        memcpy(aa, name, 5);
        aa[5] = 0;
        int res = !strcmp(aa, "upsam");
        return res;
    }
    //根据名字创建插件的方法，有两中方式:
    //一个是由权重构建
    //另一个是由序列化后的比特流创建，
    //对应了插件的两种构造函数，Upsample没有权重，对于其他有权重的插件就能够用传入的weights初始化层。mplugin是一个vector用来存储所有创建的插件层
    IPlugin *createPlugin(const char *layerName, const nvinfer1::Weights *weights, int nbWeights) override {
        assert(isPlugin(layerName));
        mPlugin.push_back(std::unique_ptr(new Upsample(2)));
        return mPlugin[mPlugin.size() - 1].get();
    }
    IPlugin *createPlugin(const char *layerName, const void *serialData, size_t serialLength) override {
        assert(isPlugin(layerName));
        
        return new Upsample(serialData, serialLength);
    }
    std::vector > mPlugin;
    //最后需要定义一个destroy方法来释放所有创建的插件层。
    void destroyPlugin() 
    {
        for (unsigned int i = 0; i < mPlugin.size(); i++) 
        {
            mPlugin[i].reset();
        }
    } 
}

对于prototxt存在多个多种插件的情况，可以在isPlugin，createPlugin方法中添加新的条件分支，根据层的名字创建对应的插件层。

实现了PluginFactory之后在调用caffeparser的时候需要设置使用它，在调用parser->parser之前加入如下代码

PluginFactory pluginFactory;
parser->setPluginFactoryExt(&pluginFactory);

就可以设置parser按照pluginFactory里面定义的规则来创建插件层，这样之前出现的不能解析Upsample层的错误就不会再出现了。

• 除了pluginExt，tensorRT中插件基类还有IPlugin，IPluginV2，继承这些基类所需要实现的类方法有细微区别，具体情况可自行查看tensorRT安装文件夹下的include/NvInfer.h文件。同时添加自己写的层到网络时的函数有addPlugin，addPluginExt，addPluginV2这几种和IPlugin，IPluginExt，IPluginV2一一对应，不能够混用，否则有些默认调用的类方法不会调用的，比如用addPlugin添加的PluginExt层是不会调用configureWithFormat方法的，因为IPlugin类没有该方法。同样的在还有caffeparser的setPluginFactory和setPluginFactoryExt也是不能混用的
• 运行程序出现cuda failure一般情况下是由于将内存数据拷贝到磁盘时出现了非法内存访问，注意检查buffer开辟的空间大小和拷贝过去数据的大小是否一致
• 有一些操作在tensorRT中不支持但是可以通过一些支持的操作进行组合替代，这样可以省去一些编写自定义层的时间
• tensorflow中的flatten操作默认时keepdims=False的，但是在转化uff文时会默认按照keepdims=True转换，因此在tensorflow中对flatten后的向量进行transpose、expanddims等等操作，在转换到uff后用tensorRT解析时容易出现错误，比如“Order size is not matching the number dimensions of TensorRT” 。最好设置tensorflow的reduce，flatten操作的keepdims=True，保持层的输出始终为4维形式，能够有效避免转到tensorRT时出现各种奇怪的错误。
• tensorRT中的slice层存在一定问题，我用network->addSlice给网络添加slice层后，在执行buildengine这一步时就会出错nvinfer1::builder::checkSanity(const nvinfer1::builder::Graph&): Assertion `tensors.size() == g.tensors.size()’ failed.，构建网络时最好避开使用slice层，或者自己实现自定层来执行slice操作。