tensorRT加速caffe模型的部署

一、简介及其tensorRT加速原理

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

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

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

  • 图 1

神经网络FPGA加速器有哪些优化方法 神经网络加速卡_神经网络FPGA加速器有哪些优化方法_02


图 2


神经网络FPGA加速器有哪些优化方法 神经网络加速卡_自定义_03


图 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函数中打印出来。