win10 使用TensorRT部署 yolov5-v4.0

辛琪

2022-01-28 2022-01-28 约 809 字预计阅读 2 分钟次阅读

作者csdn

win10 使用TensorRT部署 yolov5-v4.0（C++）

一、环境及文件准备

安装 CUDA+cudnn+TensorRT 【查看上一篇文章】
默认安装Anaconda并安装了pytorch开发环境
默认安装了 vs2019 + opencv + cmake
下载yolov5源码：https://github.com/ultralytics/yolov5/tags
下载yolov5权重：https://github.com/ultralytics/yolov5/releases
下载dirent.h：https://github.com/tronkko/dirent/blob/master/include/dirent.h 或者点击下载
下载tensorrtx：https://github.com/wang-xinyu/tensorrtx/tags （与自己训练的yolov5-xx版本一致）

二、编译

yolov5s.wts生成：将tensorrtx源码中的gen_wts.py复制到yolov5源码中并运行，生成.wts模型。
将dirent.h放置在工程目录中（随意放置）

修改CMakeLists.txt #1-10即可，参数详情查看

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71


cmake_minimum_required(VERSION 2.8)

#=========================================================
project(yolov5)                                 #1 工程名
set(OpenCV_DIR "E:\\opencv3\\opencv\\build")    #2 opencv目录
set(OpenCV_INCLUDE_DIRS ${OpenCV_DIR}\\include) #3 
set(OpenCV_LIB_DIRS ${OpenCV_DIR}\\x64\\vc15\\lib) #4
set(OpenCV_Debug_LIBS "opencv_world3412d.lib")     #5
set(OpenCV_Release_LIBS "opencv_world3412.lib")    #6
set(TRT_DIR "E:\\Downloads\\TensorRT-7.2.1.6.Windows10.x86_64.cuda-11.0.cudnn8.0\\TensorRT-7.2.1.6")  #7
set(TRT_INCLUDE_DIRS ${TRT_DIR}\\include) #8
set(TRT_LIB_DIRS ${TRT_DIR}\\lib)         #9
set(Dirent_INCLUDE_DIRS "F:\\yolov5trt")  #10
#=========================================================

add_definitions(-std=c++11)

option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)

set(THREADS_PREFER_PTHREAD_FLAG ON)
find_package(Threads)

# setup CUDA
find_package(CUDA REQUIRED)
message(STATUS "    libraries: ${CUDA_LIBRARIES}")
message(STATUS "    include path: ${CUDA_INCLUDE_DIRS}")

include_directories(${CUDA_INCLUDE_DIRS})

####
enable_language(CUDA)  # add this line, then no need to setup cuda path in vs
####
include_directories(${PROJECT_SOURCE_DIR}/include) #11
include_directories(${TRT_INCLUDE_DIRS}) #12
link_directories(${TRT_LIB_DIRS}) #13
include_directories(${OpenCV_INCLUDE_DIRS}) #14
link_directories(${OpenCV_LIB_DIRS}) #15
include_directories(${Dirent_INCLUDE_DIRS}) #16


# -D_MWAITXINTRIN_H_INCLUDED for solving error: identifier "__builtin_ia32_mwaitx" is undefined
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -D_MWAITXINTRIN_H_INCLUDED")

# setup opencv
find_package(OpenCV QUIET
    NO_MODULE
    NO_DEFAULT_PATH
    NO_CMAKE_PATH
    NO_CMAKE_ENVIRONMENT_PATH
    NO_SYSTEM_ENVIRONMENT_PATH
    NO_CMAKE_PACKAGE_REGISTRY
    NO_CMAKE_BUILDS_PATH
    NO_CMAKE_SYSTEM_PATH
    NO_CMAKE_SYSTEM_PACKAGE_REGISTRY
)

message(STATUS "OpenCV library status:")
message(STATUS "    version: ${OpenCV_VERSION}")
message(STATUS "    lib path: ${OpenCV_LIB_DIRS}")
message(STATUS "    Debug libraries: ${OpenCV_Debug_LIBS}")
message(STATUS "    Release libraries: ${OpenCV_Release_LIBS}")
message(STATUS "    include path: ${OpenCV_INCLUDE_DIRS}")

add_executable(yolov5 ${PROJECT_SOURCE_DIR}/yolov5.cpp ${PROJECT_SOURCE_DIR}/common.hpp ${PROJECT_SOURCE_DIR}/yololayer.cu ${PROJECT_SOURCE_DIR}/yololayer.h)   #17
target_link_libraries(yolov5 "nvinfer" "nvinfer_plugin") #18
target_link_libraries(yolov5 debug ${OpenCV_Debug_LIBS}) #19
target_link_libraries(yolov5 optimized ${OpenCV_Release_LIBS}) #20
target_link_libraries(yolov5 ${CUDA_LIBRARIES}) #21
target_link_libraries(yolov5 Threads::Threads)  

点击Configure，generate，Open Project
依次编译Debug-x64, Release-x64

三、测试

打开终端依次执行：（将权重文件复制到生成的目录下执行）
yolov5.exe -s yolov5s.wts yolov5.engine s
yolov5.exe -d ../../model/yolov5s.engine ../../test_tmp