TensorRT加速yoloV3(caffe版本)

加速yolov3需要处理三件事情:

1)upsample层;

2)图像的预处理;

3)加速结果的后处理;

贴出upsample层的代码:(参考:)

/******************************/
// Upsample Plugin
/******************************/
//*fuxueping
//2018-11-7
class UpsamplePlugin: public IPlugin
{
public:
    UpsamplePlugin();
    UpsamplePlugin(size_t scale)
    {
      m_scale = scale;
//      std::cout<<"UpsampleLayer0"<(data);
        const char *a = d;
        m_input_c = read(d);
        m_input_h = read(d);
        m_input_w = read(d);//input_h=input_w
        m_scale = read(d);

        m_output_count = m_input_c * m_input_w * m_input_h * m_scale * m_scale;
        assert(d == a + length);
    }

    ~UpsamplePlugin();

    int getNbOutputs() const override{
      return 1;
    }

    Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override{
        m_input_c=inputs[0].d[0];
        m_input_h=inputs[0].d[1];
        m_input_w=inputs[0].d[2];

        assert(index == 0 && nbInputDims == 1 && inputs[0].nbDims == 3);
        return DimsCHW(inputs[0].d[0], inputs[0].d[1]*m_scale, inputs[0].d[2]*m_scale);
    }

    void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override{
        m_input_c = inputDims[0].d[0];
        m_input_h = inputDims[0].d[1];
        m_input_w = inputDims[0].d[2];

     }

    int initialize() override{
        return 0;
    }

    void terminate() override{

    }
    size_t getWorkspaceSize(int maxBatchSize) const override{
      return 0;
    }

    int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override{
        int d1 = m_input_c;
        int d2 = m_input_h*m_scale;
        int d3 = m_input_w*m_scale;
//        cout<<"d1 = "<if (layer_type == "Upsample")
                    {
                        int scale = iter1->second.upsample_param.scale;
                        m_upsample_plugin_ptr = std::unique_ptr(new UpsamplePlugin(scale));
                        m_vec_upsample_plugin.push_back(std::move(m_upsample_plugin_ptr));
                        return m_vec_upsample_plugin.back().get();
//                        _nvPlugins[layerName] = (plugin::INvPlugin*)(new UpsamplePlugin(scale));
//                        return _nvPlugins.at(layerName);
                    }

图片的前向处理和后期结果处理,参考的demo;(地址:https://download.csdn.net/download/qq_22764813/10798442)

还需要使用代码将darknet的yoloV3转换为caffe版本,这个我就帮不了大家了,需要自己处理.如果有疑问或者什么错误,希望大家留言,谢谢;

 

 

你可能感兴趣的:(TensorRT)