caffe-ssd网络模型 tensorRT加速

tensorRT加速caffe模型的部署

一、简介及其tensorRT加速原理

在计算资源并不丰富的嵌入式设备上，TensorRT之所以能加速神经网络的的推断主要得益于两点。首先是TensorRT支持INT8和FP16的计算，通过在减少计算量和保持精度之间达到一个理想的trade-off，达到加速推断的目的。

更为重要的是TensorRT对于网络结构进行了重构和优化，主要体现在一下几个方面。

tensorRT通过解析网络模型将网络中无用的输出层消除以减小计算。
对于网络结构的垂直整合，即将目前主流神经网络的conv、BN、Relu三个层融合为了一个层，例如将图1所示的常见的Inception结构重构为图2所示的网络结构。
对于网络的水平组合，水平组合是指将输入为相同张量和执行相同操作的层融合一起，如图2向图3的转化。
对于concat层，将contact层的输入直接送入下面的操作中，不用单独进行concat后在输入计算，相当于减少了一次传输吞吐。

图 1

caffe-ssd网络模型 tensorRT加速

图 2

caffe-ssd网络模型 tensorRT加速

图 3

二、改写deploy.prototxt

将deploy.prototxt文件名改写为deploy_plugin.prototxt

convolution层的param{}全部去掉，convolution_param中的weight_filter{}去掉，bias_filter{}去掉
将自定义层的名字改写为IPlugin，自定义层的参数不用写在prototxt中，全部写在IPlugin的实现中。例如SSD的priorbox层，permute层，都是需要改写为IPlugin
SSD的最后输出是detection_out层，需要将detection_out_param去掉，然后新增加一个输出detection_out2,因为在tensorRT中默认的输出是两个，如果只有一个top，那么程序会报错，“Plugin layer output count is not equal to caffe output count“

三、完成自定义层的代码

先实现PluginFactory的class，继承IPluginFactory，class中定义需要使用的自定义层。

完善判断函数

//将所有的plugin层的名字全部写到strcmp()中
//通过此函数来判断是否是自定义的IPlugin层
bool PluginFactory::isPlugin(const char* name)
{
    return (!strcmp(name,"conv4_3_norm")     ||!strcmp(name,"conv4_3_norm_mbox_conf_perm"));
}

完善创建plugin函数

//
nvinfer1::IPlugin* PluginFactory::createPlugin(const char* layerName,
                                              const nvinfer1::Weights* weights,
                                              int nbWeights)
{
    assert(isPlugin(layerName));
    //通过if语句判断是哪一个plugin，然后实现相对应的功能
    //下面这个是normalize层的实现，使用tensorRT中自带的createSSDNormalizePlugin函数来实现
    if(!strcmp(layerName,"conv4_3_norm"))
    {
        assert(mNormalizerLayer.get() == nullptr);
        mNormalizerLayer = std::unique_ptr<INvPlugin,decltpe(nvPluginDeleter)>(createSSDNormalizePlugin(weights,false,false,0.0001),nvPluginDeleter);
        return mNormalizeLayer.get();
    }
    //下面的这个SSD的permute层
    //0,2,3,1的数据就是deploy.prototxt中的conv4_3_norm_mbox_conf_perm这一层的数据
    else if(!strcmp(layerName,"conv4_3_norm_mbox_conf_perm"))
    {
        assert(mConv4_3_norm_mbox_conf_perm_layer.get() == nullptr);
        mConv4_3_norm_mbox_conf_perm_layer = std::unique_ptr<INvPlugin,decltype(nvPluginDelter)>(createSSDPermutePlugin({{0,2,3,1}}),nvPluginDelter);
        return mConv4_3_norm_mbox_conf_perm_layer.get();
    }
    //下面这个是priorbox层
    //priorbox层的每一个参数都对应于deploy.prototxt中的参数
    //需要仔细对比deploy.prototxt文件
    else if(!strcmp(layerName,"pool6_mbox_priorbox"))
    {
        assert(mPool6_mbox_priorbox_layer.get() == nullptr);
        PriorBoxParameters params;
        float minsize[1] = {200};
        float maxsize[1] = {240};
        float aspect_ratio[2] = {1.0,2.0};
        params.minSize = minsize;
        params.maxSize = maxsize;
        params.aspectRatio = aspect_raio;
        params.numMinSize = 1;
        params.numAspectRatios = 2;
        params.numMaxSize = 1;
        params.flip = true;
        params.clip = false;
        params.variance[0] = 0.1;
        params.variance[1] = 0.1;
        params.variance[2] = 0.2;
        params.variance[3] = 0.2;
        params.imgH = 0;
        params.imgW = 0;
        params.stepH = 300;
        params.stepW = 300;
        params.offset = 0.5;
        mPool6_mbox_priorbox_layer = std::unique_ptr<INvPlugin,decltype(nvPluginDeleter)>(createSSDPriorBoxPlugin(params),nvPluginDeleter);
        return mPool6_mbox_priorbox_layer.get();
    }
}

如果tensorRT中没有对应的实现，那就需要自己手撸代码了。例如SSD中的softmax(axive = 2)这种情况，tensorRT中原生不支持，需要自己用代码实现

//softmax layer
else if(!strcmp(layerName,"mbox_conf_softmax"))
{
    assert(mPluginSoftmax == nullptr);
    assert(nbWeights == 0 && weights == nullptr);
    mPluginSoftmax = std::unique_ptr<SoftmaxPlugin>(new SoftmaxPlugin());
    return mPluginSoftmax.get();
}

//手动实现softmax
class SoftmaxPlugin : public IPlugin
{
public:
	int initialize() override {return 0;}
	inline void terminate() override {}
	SoftmaxPlugin(){}
	SoftmaxPlugin(const void* buffer,size_t size)
    {
    	assert(size == sizeof(mCopySize));
    	mCopySize = *reinterpret_cast<const size_t*>(buffer);
    }
    
    inline int getNbOutputs() const override
    {
    	return 1;
    }
    
    Dims getOutputDimensions(int index,const Dims* inputs,int nbInputDims) override
    {
    	assert(nbInputDims == 1);
    	assert(index == 0);
    	assert(inputs[index].nbDims == 3);
    	//这里的5是最后检测的类别数
    	assert((inputs[0].d[0]) * (inputs[0].d[1] % 5 == 0);
    	
    	return DimsCHW(inputs[0].d[0] , inputs[0].d[1],inputs[0].d[2]);
    }
    
    size_t getWorkspaceSize(int) const override
    {
    	return nCopySize * 1;
    }
    
    //最重要的是这里的函数
    int enqueue(int batchSize,const void*const *input,void** outputs,void* workspace,cudaStream_t stream) override
    {
    	cudaSoftmax(1917*5,5,(float*)*inputs,static_cast<float *>(*outputs));
    	return 0;
    }
    
}

cudaSoftmax函数用cuda代码实现，查看.cu文件

tensorRT的使用

caffeToTRTModel

void TensorNet::caffeToTRTModel(const std::string& deployFile, const std::string& modelFile, const std::vector<std::string>& outputs,
                                unsigned int maxBatchSize)
{
    IBuilder* builder = createInferBuilder(gLogger);
    INetworkDefinition* network = builder->createNetwork();

    ICaffeParser* parser = createCaffeParser();
    parser->setPluginFactory(&pluginFactory);

    bool useFp16 = builder->platformHasFastFp16();
    //useFp16 = false;
    DataType modelDataType = useFp16 ? DataType::kHALF : DataType::kFLOAT;

    std::cout << deployFile.c_str() <<std::endl;
    std::cout << modelFile.c_str() <<std::endl;
    //std::cout << (*network) <<std::endl;
    std::cout << "Here : 1"<<std::endl;
    const IBlobNameToTensor* blobNameToTensor =	parser->parse(deployFile.c_str(),modelFile.c_str(),*network,                DataType::kFLOAT);
    std::cout << "Here : 2" <<std::endl;
    assert(blobNameToTensor != nullptr);
    std::cout << "Here : 3" <<std::endl;
    for (auto& s : outputs) network->markOutput(*blobNameToTensor->find(s.c_str()));

    builder->setMaxBatchSize(maxBatchSize);
    builder->setMaxWorkspaceSize(10 << 20);
    std::cout << "Here : 4"<< std::endl;
    ICudaEngine* engine = builder->buildCudaEngine( *network );
    std::cout << "Here : 5"<<std::endl;
    assert(engine);
    network->destroy();
    parser->destroy();
    gieModelStream = engine->serialize();
    engine->destroy();
    builder->destroy();
    pluginFactory.destroyPlugin();
    shutdownProtobufLibrary();
}

这里的函数中

bool useFp16 = builder->platformHasFastFp16();

显卡是不支持fp16的计算的，但是tx2等嵌入式平台支持fp16，调用此函数可以得到是否支持fp16.

四、需要注意的地方

修改prototxt文件一定要仔细，确认没有prototxt的语法错误，然后根据IPlugin的数量来改写代码，IPlugin的数量和isPlugin函数的数量保持一致
priorbox层中的参数，按照原始的deploy改写的时候需要注意,在deploy.prototxt中的priorbox的参数，例如aspect_ratio: 2.0,则在代码中需要写成float aspect_ratio[2] = {1.0,2.0};如果是aspect_ratio: 3.0,则需要写成`float aspect_ratio[3] = {1.0,2.0,3.0};
detection_out层中的参数，tensorRT的版本不同，命名规则是不一样的，需要根据tensorRT版本来改变。
在softmax的参数，需要有priorbox的数量，例如vgg6-ssd的数量的8732，resnet20-ssd的数量是1917，红绿灯ssd的数量是4992。这个参数可以在reshape层中打印出来看到。在reshape类定义中，getOutputDimensions函数中打印出来。

秒客网

caffe-ssd网络模型 tensorRT加速

tensorRT加速caffe模型的部署

相关文章