在Jetson Nano上使用TensorRT C++实现YOLOv5模型推理

软件版本信息

软件版本信息

 

JetPack4.6
CUDA10.2
TensorRT8.0.1
OpenCV4.5.4

 

导出YOLOv5 engine文件

前面有一篇文章详细说明了如何在Jetson Nano上安装YOLOv5，然后运行，这里只需在导出的时候导出engine模型文件支持，命令行如下：

python export.py --weights yolov5.pt --include onnx engine

这里需要注意的TensorRT版本一致问题。如果engine文件不是在Jetson Nano上生成的，而在其他PC机器上生成，则TensorRT版本必须与Jetson Nano上使用的版本保持一致。

 

TensorRT推理

首先创建编译CMakeLists.txt文件，然后把下面的内容copy进去：

cmake_minimum_required( VERSION 2.8 )
# 声明一个 cmake 工程
project(yolov5_tensorrt_demo)
# 设置编译模式
#set( CMAKE_BUILD_TYPE "Release" )
#添加OPENCV库
#指定OpenCV版本，代码如下
#find_package(OpenCV 4.5.4 REQUIRED)
#如果不需要指定OpenCV版本，代码如下
find_package(OpenCV REQUIRED)
find_package(CUDA REQUIRED)


include_directories(
       ./src/)
#添加OpenCV头文件
include_directories(${OpenCV_INCLUDE_DIRS})
# 添加CUDA10.2头文件
include_directories(/usr/local/cuda-10.2/include)
link_directories(/usr/local/cuda-10.2/lib64)
# tensorRT
include_directories(/usr/local/cuda/include)
link_directories(/usr/lib/arrch64-linux-gnu)
#显示OpenCV_INCLUDE_DIRS的值
message(${OpenCV_INCLUDE_DIRS})


FILE(GLOB_RECURSE TEST_SRC
       src/*.cpp
       )
# 添加一个可执行程序
# 语法：add_executable( 程序名 源代码文件 ）
add_executable(target main.cpp ${TEST_SRC})
# 将库文件链接到可执行程序上
target_link_libraries(target  nvinfer)
target_link_libraries(target  cudart)
target_link_libraries(target  ${OpenCV_LIBS})

 

构建的目录结构如下

然后在Jetson Nano上打开终端命令行输入：

 

mkdir build
cd build
cmake ..

 

运行结果如下：

然后运行make

成功运行之后，执行编译生成的target文件：

 

./target

 

运行推理与显示界面如下：

错误与修正

01错误一：

error: 'cudaMalloc' was not declared in this scope

这个是因为没有添加下面的头文件：

 

#include 
#include 

 

添加之后就可以解决。

02错误二：

Engine文件版本一致导致，原因是我之前导出的是在tensorRT8.4版本，jetsonNano是8.0的，重新在Jetson Nano上导出一下就可以解决。   最后贴一下，演示的程序main C++代码：

#include 
#include 
#include 
#include 

#include "tensorrt_yolov5_demo.h"

using namespace cv;

std::string label_map = "classes.txt";
int main(int argc, char** argv) {
    std::vector classNames;
    std::ifstream fp(label_map);
    std::string name;
    while (!fp.eof()) {
        getline(fp, name);
        if (name.length()) {
            classNames.push_back(name);
        }
    }
    fp.close();

    auto detector = std::make_shared();
    detector->initConfig("yolov5s.engine", 0.25, 0.25);
    std::vector results;
    cv::VideoCapture capture("example_dsh.mp4");
    cv::Mat frame;
    while (true) {
        bool ret = capture.read(frame);
        if (!ret) {
            break;
        }
        detector->detect(frame, results);
        for (DetectResult dr : results) {
            cv::Rect box = dr.box;
            cv::putText(frame, classNames[dr.classId], cv::Point(box.tl().x, box.tl().y - 10), cv::FONT_HERSHEY_SIMPLEX, .5, cv::Scalar(0, 0, 0));
        }
        cv::imshow("YOLOv5-6.1 + TensorRT8.4 + Jetson Nano - by gloomyfish", frame);
        char c = cv::waitKey(1);
        if (c == 27) { // ESC 退出
            break;
        }
        // reset for next frame
        results.clear();
    }
    cv::waitKey(0);
    cv::destroyAllWindows();
    return 0;
}

 

审核编辑：汤梓红