mnist手写数字识别是一个非常经典的模型了,体积小,结构简单,特别适合作为入门案例。参考了官方例子,我根据自己的使用习惯进行了梳理。
运行环境:
- Windows10
- Visual Studio 2015 Update3
- Cuda 10.0
- cuDNN 7.6.5
- TensorRT 7.0.0.11
- OpenCV 4.3.0
1.TensorRT推理的12个基本步骤
step1:创建runtime
step2:反序列化创建engine
step3:创建context
step4:获取输入输出索引
step5:创建buffers
step6:为输入输出开辟GPU显存
step7:创建cuda流
step8:从CPU到GPU----拷贝input数据
step9:异步推理
step10:从GPU到CPU----拷贝output数据
step11:同步cuda流
step12:释放资源
下面我就通过mnist的案例来逐步实现。
2.TensorRT运行库导入
TensorRT头文件
#include “NvInfer.h”
#include “NvCaffeParser.h”
#include “NvOnnxParser.h”
#include “NvUffParser.h”
#include “common.h”
using namespace nvinfer1;
using namespace nvcaffeparser1;
using namespace nvuffparser;
TensorRT库
cudart.lib
nvinfer.lib
nvinfer_plugin.lib
myelin64_1.lib
nvonnxparser.lib
nvparsers.lib
OpenCV头文件
#include <opencv2/opencv.hpp>
OpenCV库
opencv_world430.lib
3.下面推理执行过程
第一步:设置参数
为了方便使用,我定义了两个存储结构
struct TensorRT {
IExecutionContext* context;
ICudaEngine* engine;
IRuntime* runtime;
};
struct TRT_CONFIG
{
std::string networkName; //网络名称
std::string trtModelName; //trt模型名称
std::string onnxModeName; //onnx模型名称
std::string input_blob_name; //输入网络名
std::string output_blob_name; //输出网络名
std::string output_config_name; //输出配置名
int batchSize; //batch size
int dlaCore; //DLA是否打开
int runInFp16; //是否采用FP16
int runInInt8; //是否采用FP8
int input_h; //输入图像高度
int input_w; //输入 图像宽度
int input_c; //输入图像通道数
int output_size; //输出size
int class_num; //输出类别数
int output_len; //输出长度
};
//定义配置
TRT_CONFIG trtConfig;
trtConfig.networkName = "TRT.Mnist";
trtConfig.onnxModeName = "mnist.onnx";
trtConfig.trtModelName = "mnist.trt";
trtConfig.input_blob_name = "Input3";
trtConfig.output_blob_name = "Plus214_Output_0";
trtConfig.batchSize = 1;
trtConfig.dlaCore =1;
trtConfig.runInFp16 = 1;
trtConfig.runInInt8 = 0;
trtConfig.input_h = 28;
trtConfig.input_w = 28;
trtConfig.input_c = 1;
trtConfig.output_size = 10;
m_rtInterface.InitParams(trtConfig);
第二步:加载模型
IHostMemory* trtModelStream{
nullptr };
TensorRT* ptensor_rt;
IExecutionContext* context = nullptr;
IRuntime* runtime = nullptr;
ICudaEngine* engine = nullptr;
//这里我们直接加载TRT模型,省略了模型转换步骤,后面再讲
ptensor_rt = m_rtInterface.LoadTRTModel(trtConfig.trtModelName.c_str())