trt_inference.h

跳转至此文件的文档。

/*
 * Copyright (c) 2016, NVIDIA CORPORATION。保留所有权利。
 *
 * 源代码和二进制形式的再分发和使用，无论是否经过
 * 修改，均被允许，前提是满足以下条件
 * ：
 * * 源代码的再分发必须保留上述版权声明
 * 、此条件列表以及以下免责声明。
 * * 二进制形式的再分发必须在上述版权声明中复制
 * 、此条件列表以及以下免责声明，位于
 * 文档和/或随发行版提供的其他材料中。
 * * NVIDIA CORPORATION 的名称及其
 * 贡献者的名称不得用于认可或推广衍生自
 * 本软件的产品，除非获得明确的事先书面许可。
 *
 * 本软件由版权所有者``按原样''提供，并且不作任何
 * 明示或暗示的保证，包括但不限于
 * 对适销性和特定用途适用性的暗示保证。
 * 目的不作担保。在任何情况下，版权所有者或
 * 贡献者均不对任何直接、间接、偶然、特殊、
 * 惩戒性或后果性损害（包括但不限于
 * 采购替代商品或服务； 无法使用、数据或
 * 利润； 或业务中断）承担责任，无论其成因和任何责任理论如何
 * 的责任，无论是合同、严格责任还是侵权
 * （包括疏忽或其他）因使用而以任何方式引起的
 * 本软件，即使已被告知可能发生此类损害。
 */
#ifndef TRT_INFERENCE_H_
#define TRT_INFERENCE_H_

#include <fstream>
#include <queue>
#include "NvInfer.h"
#include "NvCaffeParser.h"
#include "NvOnnxParser.h"
#include "opencv2/video/tracking.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/highgui/highgui.hpp"
#include <opencv2/objdetect/objdetect.hpp>
using namespace nvinfer1;
using namespace nvcaffeparser1;
using namespace nvonnxparser;
using namespace std;

// 模型索引
#define GOOGLENET_SINGLE_CLASS 0
#define GOOGLENET_THREE_CLASS 1
#define RESNET_THREE_CLASS 2

class Logger;

class Profiler;

class TRT_Context
{
public
 // 网络相关参数
 int getNetWidth() const;

 int getNetHeight() const;

 uint32_t getBatchSize() const;

 int getChannel() const;

 int getModelClassCnt() const;

 void* getScales() const;

 void* getOffsets() const;

 // 缓冲区在 TRT_Context 中分配，
 // 暴露此接口以输入数据
 void*& getBuffer(const int& index);

 float*& getInputBuf();

 uint32_t getNumTrtInstances() const;

 //0 fp16 1 fp32 2 int8
 void setMode(const int& mode);

 void setBatchSize(const uint32_t& batchsize);

 void setDumpResult(const bool& dump_result);

 void setTrtProfilerEnabled(const bool& enable_trt_profiler);

 int getFilterNum() const;
 void setFilterNum(const unsigned int& filter_num);

 TRT_Context();

 void setModelIndex(int modelIndex);

 void buildTrtContext(const string& deployfile,
 const string& modelfile, bool bUseCPUBuf = false, bool isOnnxModel = false);

 void doInference(
 queue< vector<cv::Rect> >* rectList_queue,
 float *input = NULL);

 void destroyTrtContext(bool bUseCPUBuf = false);

 ~TRT_Context();

private
 int net_width;
 int net_height;
 int filter_num;
 void **buffers;
 float *input_buf;
 float *output_cov_buf;
 float *output_bbox_buf;
 void* offset_gpu;
 void* scales_gpu;
 float helnet_scale[4];
 IRuntime *runtime;
 ICudaEngine *engine;
 IExecutionContext *context;
 uint32_t *pResultArray;
 int channel; //输入文件通道数
 int num_bindings;
 int trtinstance_num; //推理通道数
 int batch_size;
 int mode;
 bool dump_result;
 ofstream fstream;
 bool enable_trt_profiler;
 bool is_onnx_model;
 IHostMemory *trtModelStream{nullptr};
 vector<string> outputs;
 string result_file;
 Logger *pLogger;
 Profiler *pProfiler;
 int frame_num;
 uint64_t elapsed_frame_num;
 uint64_t elapsed_time;
 int inputIndex;
 int outputIndex;
 int outputIndexBBOX;
 Dims3 inputDims;
 Dims3 outputDims;
 Dims3 outputDimsBBOX;
 size_t inputSize;
 size_t outputSize;
 size_t outputSizeBBOX;

 struct {
 const int classCnt;
 float THRESHOLD[3];
 const char *INPUT_BLOB_NAME;
 const char *OUTPUT_BLOB_NAME;
 const char *OUTPUT_BBOX_NAME;
 const int STRIDE;
 const int WORKSPACE_SIZE;
 int offsets[3];
 float input_scale[3];
 float bbox_output_scales[4];
 const int ParseFunc_ID;
 } *g_pModelNetAttr, gModelNetAttr[4] = {
 {
 // GOOGLENET_SINGLE_CLASS
 1,
 {0.8, 0, 0},
 "data",
 "coverage",
 "bboxes",
 4,
 450 * 1024 * 1024,
 {0, 0, 0},
 {1.0f, 1.0f, 1.0f},
 {1, 1, 1, 1},
 0
 },

 {
 // GOOGLENET_THREE_CLASS
 3,
 {0.6, 0.6, 1.0}, //人, 摩托车, 汽车
 "data",
 "Layer16_cov",
 "Layer16_bbox",
 16,
 110 * 1024 * 1024,
 {124, 117, 104},
 {1.0f, 1.0f, 1.0f},
 {-640, -368, 640, 368},
 0
 },

 {
 // RESNET_THREE_CLASS
 4,
 {0.1, 0.1, 0.1}, //人, 摩托车, 汽车
 "data",
 "Layer7_cov",
 "Layer7_bbox",
 16,
 110 * 1024 * 1024,
 {0, 0, 0},
 {0.0039215697906911373, 0.0039215697906911373, 0.0039215697906911373},
 {-640, -368, 640, 368},
 1
 },
 };
 enum Mode_type{
 MODE_FP16 = 0,
 MODE_FP32 = 1,
 MODE_INT8 = 2
 };
 int parseNet(const string& deployfile);
 void parseBbox(vector<cv::Rect>* rectList, int batch_th);
 void ParseResnet10Bbox(vector<cv::Rect>* rectList, int batch_th);
 void allocateMemory(bool bUseCPUBuf);
 void releaseMemory(bool bUseCPUBuf);
 void caffeToTRTModel(const string& deployfile, const string& modelfile);
 void onnxToTRTModel(const string& modelfile);
};

#endif