tensorrt nvcaffeparser1::IPluginFactoryV2

私有成员

一些需要传入层接口的参数除了(weights,bias)

PluginFactory(TrtPluginParams params)

给PluginFactory私有成员初始化

PluginFactory::PluginFactory(TrtPluginParams params) {
    spdlog::info("create plugin factory");mYoloClassNum = params.yoloClassNum;mYolo3NetSize = params.yolo3NetSize;mUpsampleScale = params.upsampleScale;spdlog::info("yolo3 params: class: {}, netSize: {} ",mYoloClassNum,mYolo3NetSize);spdlog::info("upsample params: scale: {}",mUpsampleScale);
}

virtual ~PluginFactory() {}

重写为空值

virtual bool isPluginV2(const char* layerName) override;

判断层名是否在plugins之中

bool PluginFactory::isPluginV2(const char* layerName) 
{
    std::string strName{
    layerName};std::transform(strName.begin(), strName.end(), strName.begin(), ::tolower);return (strName.find("prelu") != std::string::npos || strName.find("upsample") != std::string::npos|| strName.find("yolo-det") != std::string::npos);
}

virtual IPluginV2* createPlugin(const char* layerName, const Weights* weights, int nbWeights, const char* libNamespace="") override;

IPluginV2* PluginFactory::createPlugin(const char *layerName, const Weights* weights, int nbWeights, const char* libNamespace) 
{
    assert(isPluginV2(layerName));std::string strName{
    layerName};std::transform(strName.begin(), strName.end(), strName.begin(), ::tolower);if (strName.find("prelu") != std::string::npos) {
    // std::cout << "nbWeight: " << nbWeights << std::endl;// std::cout << "weights.count: " << weights->count << std::endl;return (IPluginV2*)(new PReLUPlugin(weights, nbWeights));} else if(strName.find("upsample") != std::string::npos) {
    return (IPluginV2*)(new UpSamplePlugin(mUpsampleScale));}else if(strName.find("yolo-det") != std::string::npos) {
    return (IPluginV2*)(new YoloLayerPlugin(mYoloClassNum,mYolo3NetSize));}else{
    std::cout << "warning : " << layerName << std::endl;assert(0);return nullptr;}
}