- 把
TensorRTC++api推理YOLOv5的代码，打包成动态链接库，并通过Python调用。
这样做的原因：
-
使用
TensorRT的C++api优化YOLO模型，可以显著提升目标检测的速度，但C++代码不方便扩展为网络通信协议接口； -
Python作为最流行的胶水语言，拥有很多成熟的通信协议库，可以方便的进行各种网络协议通信，赋能于各种各样的服务；
因此，若将C++的模型推理代码编译为动态链接库，再使用Python封装，那就既有推理速度的优势，又具备可扩展性的便利
可能有人问为什么不直接使用 TensorRT Python Api，看似可以达到相同的效果，作者认为本项目有以下优势，比直接使用TensorRT Python Api 更好：
- 使用
TensorRT Python Api时，前处理和后处理需使用Python完成，不过也是通过numpy实现的，速度方面没有太大差距； - 作者使用 Cuda 编程加速预处理，提升预处理速度，之后一同打包到动态链接库，推理速度更快。
- 使用
TensorRT-v8.2的C++api，加速YOLOv5-v5.0目标检测； - 在
Linux x86_64上进行部署； - 在
Jetson系列嵌入式设备上也是可行的，把本项目中的CMakeLists.txt文件中头文件、库文件相关目录更换即可；
大致实现过程如下：
- 作者把
YOLOv5的TensorRT推理封装成C++的类，关键代码如下：
class YoloDetecter
{
public:
YoloDetecter(const std::string trtFile, const int gpuId);
~YoloDetecter();
float* inference(cv::Mat& img);
};- 接着作者使用
C类型的函数再次封装上面的类，关键代码如下：
#ifdef __cplusplus
extern "C" {
#endif
YoloDetecter* YoloDetecter_new(char* trtFile, int gpuId){
return new YoloDetecter(std::string(trtFile), gpuId);
}
float* inference_one(YoloDetecter* instance, const uchar* srcImgData, const int srcH, const int srcW){
cv::Mat srcImg(srcH, srcW, CV_8UC3);
memcpy(srcImg.data, srcImgData, srcH * srcW * 3 * sizeof(uchar));
return instance->inference(srcImg);
}
void destroy(YoloDetecter* instance) { delete instance; }
#ifdef __cplusplus
}
#endif- 把封装后的代码生成动态链接库，
CMakeLists.txt中的关键部分如下
# ====== yolo infer shared lib ======
cuda_add_library(yolo_infer SHARED
${PROJECT_SOURCE_DIR}/src/preprocess.cu
${PROJECT_SOURCE_DIR}/src/yololayer.cu
${PROJECT_SOURCE_DIR}/src/yolo_infer.cpp
${PROJECT_SOURCE_DIR}/main.cpp
)
target_link_libraries(yolo_infer nvinfer cudart ${OpenCV_LIBS})- 再接着作者使用
Python封装一个检测类，类当中调用的是上述C/C++代码，关键部分代码如下：
class YoloDetector:
def __init__(self, trt_file, gpu_id=0):
self.yolo_infer_lib = ctypes.cdll.LoadLibrary("./lib/libyolo_infer.so")
self.cpp_yolo_detector = self.yolo_infer_lib.YoloDetecter_new(trt_file.encode('utf-8'), gpu_id)
def release(self):
self.yolo_infer_lib.destroy(self.cpp_yolo_detector)
def infer(self, image):
out_data = self.yolo_infer_lib.inference_one(self.cpp_yolo_detector, image, height, width)
out_data = as_array(out_data).copy().reshape(-1)- 最后，使用者不必关心具体的实现，仅仅使用下面的 2 行代码，即可实现
Python对YOLOv5+TensorRTC++代码的调用
# 实例化目标检测类
yolo_infer = YoloDetector(trt_file=plan_path, gpu_id=0)
# 使用目标检测实例推理
detect_res = yolo_infer.infer(img)
- 作者自己所使用的基本环境如下：
| Ubuntu | CUDA | cuDNN | TensorRT | OpenCV |
|---|---|---|---|---|
| 20.04 | 11.6 | 8.4 | 8.2.4 | 4.5.0 |
想要方便点的话，可以直接拉取一个 docker 镜像：
docker pull nvcr.io/nvidia/tensorrt:22.04-py3然后在镜像中编译安装 OpenCV，具体可参考下面链接中的环境构建部分：
https://github.com/emptysoal/TensorRT-v8-YOLOv5-v5.0
- python 第三方库环境
pip install numpy==1.22.3
pip install opencv-python==3.4.16.59把 YOLO检测模型，转换成TensorRT的序列化文件，后缀 .plan（作者的习惯，也可以是.engine或其他）
- 链接：https://pan.baidu.com/s/1YG-A8dXL4zWvecsD6mW2ug
- 提取码：y2oz
下载并解压后，模型文件说明：
模型文件目录
└── YOLOv5-v5.0 # 该目录中存放的是 YOLOv5 目标检测网络的模型
├── yolov5s.pt # 官方 PyTorch 格式的模型文件
└── para.wts # 根据 yolov5s.pt 导出的 wts 格式模型文件也可以直接从官方YOLOv5-v5.0处下载 yolov5s.pt，然后直接进入到下面的模型转换
-
将上述
yolov5s.pt转为model.plan，或para.wts转为model.plan -
具体转换方法参考下面链接，也是作者自己发布的一个项目
-
https://github.com/emptysoal/TensorRT-v8-YOLOv5-v5.0/tree/main
完成之后便可得到 model.plan ，为检测网络的 TensorRT 序列化模型文件。
- 开始编译并运行
- 按如下步骤
# 创建用于存储 TensoRT 模型的目录
mkdir resources
# 把上面转换得到的 plan 文件复制到目录 resources 中
cp {TensorRT-v8-YOLOv5-v5.0}/model.plan ./resources
mkdir images # 向其中放入用于推理的图片文件
mkdir build
cd build
cmake ..
make
# 以上执行完成后，会生成 lib 目录，其中存放着 C++ 代码所生成的动态链接库
# 运行 python 推理代码，即可调用动态链接库完成快速推理
python main.py
# 检测结果会保存到output目录下- 运行后可以看到如下日志：
Succeeded getting serialized engine!
Succeeded loading engine!
Model load cost: 1.2174 s
Infer 001.jpg cost: 0.0077 s
Infer 002.jpg cost: 0.0054 s
Infer 003.jpg cost: 0.0043 s在Python如何调用C++代码部分，主要参考了下面的链接：
- https://geek-docs.com/python/python-ask-answer/601_python_how_can_i_use_c_class_in_python.html
- https://blog.csdn.net/qq_41554005/article/details/128292116
作者的一些其他项目，若感兴趣，欢迎交流