TensorRT：2.实现Mnist手写数字识别推理

最新推荐文章于 2024-07-06 17:02:11 发布

代丹

最新推荐文章于 2024-07-06 17:02:11 发布

阅读量1k

点赞数 3

分类专栏：人工智能文章标签： tensorrt 深度学习 mnist 模型加速

本文链接：https://blog.csdn.net/mist99/article/details/119885372

版权

本文介绍了使用TensorRT进行MNIST手写数字识别的详细步骤，包括设置参数、加载模型、图片预处理、推理预测和结果处理。环境配置为Windows10、TensorRT 7.0.0.11等。实验结果显示推理速度为1.0583ms，准确率为94.27%。

摘要由CSDN通过智能技术生成

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())