tensorrt nvcaffeparser1::IPluginFactoryV2

class PluginFactory : public 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;
    }
}