前言
刚开工 最近因有在windows端部署模型的需求
很早就听说TensorRT可以加速模型推理,但一直没时间去进行实践,今天就来把这个陈年旧坑填补一下。
本机实验环境:
window 11、rtx3060、tensorrt8.5.1.7、cuda11.7、cudnn8.8.1.3,cmake,vs2022;
因为最近刚开始写,代码会在过段时间进行上传,如果读者太想进步了,可以跟我私聊!!!!
其他的一些东西大家应该都会,那么话不多说直接上干货!!!!
模型转换
采用的是pt->onnx->engine路线,使用yolov5官方仓库的exoprt和tensorrt的trtexec直接进行转换。
本文以单batch为例输入是[1,3,640,640],输出是[1,25200,85],若是需要多batch可以自行修改export和trtexec的转换命令。
模型部署
1.创建抽象基类 BaseInfer
和一个结构体 Detection
。
-
命名空间
命名空间用于将相关的类、函数、变量组织在一起,以避免命名冲突。在这里,TRT
:TRT
是指 TensorRT。 - 结构体
Detection
结构体定义了一个检测结果,包括:bbox[4]
:表示边界框的坐标,通常是左上角和右下角的像素坐标。classID
:表示检测到的对象类别的标识。conf
:表示检测到的对象的置信度(confidence)。- 类
BaseInfer
: BaseInfer
是一个抽象基类(abstract base class),包含了纯虚函数(pure virtual functions),需要子类实现具体的功能。这些纯虚函数包括:engineModelInit()
:用于初始化引擎模型。getBatchImagesData()
:获取批量图像数据。preprocessImages()
:图像预处理。
2.写各个函数的实现
Detinfer.h中的派生类继承自抽象类,并在Detinfer.cpp中给出推理视频流和图像的实现;
Detinfer.h
Detinfer.cpp
前后处理的cuda加速
因为输入图像的尺寸要和训练时候输入网络的图像尺寸是一样的,所以要使用gpu来加速图像的resize操作。
网络模型输出的张量的后处理,涉及到图像框的解码和nms操作:
配置文件
项目目录结构
编译
剩下的就是使用cmake来进行项目编译的工作,配置好相关路径编译后就可以进行部署了。
推理结果:本机推理一帧在2-3ms。