文章目录
- 一. yolov5 v6.0训练模型
- 二.训练好的yolov5模型转tensorrt引擎
一. yolov5 v6.0训练模型
-
官网下载yolov5 v6.0代码
-
下载官方预训练好的模型
-
安装yolov5所需要的库文件,requirements.txt在下载好的yolov5源代码中有
pip install -r C:\Users\10001540\Downloads\yolov5-6.0\requirements.txt
-
打开yolov5源代码中的detect.py文件,修改模型的位置
运行后可能出现各种错误,可以去参考网上的教程
二.训练好的yolov5模型转tensorrt引擎
-
去tensorrtx官网下载代码
-
将tensorrtx下的yolov5中的gen_wts.py复制到yolov5源代码文件夹中
-
参考yolov5官方说明,将yolov5模型文件yolov5s.pt转换为yolov5s.wts文件
python gen_wts.py -w weights/yolov5s.pt -o yolov5s.wts
-
进入tensorrtx下的yolov5文件夹,修改里面的CMakeList.txt,如下:
cmake_minimum_required(VERSION 2.6)project(yolov5) #1 set(OpenCV_DIR "D:\\Program Files\\opencv\\build") #2 set(OpenCV_INCLUDE_DIRS ${OpenCV_DIR}\\include) #3 set(OpenCV_LIB_DIRS ${OpenCV_DIR}\\x64\\vc15\\lib) #4 set(OpenCV_Debug_LIBS "opencv_world454d.lib") #5 set(OpenCV_Release_LIBS "opencv_world454.lib") #6 set(TRT_DIR "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\TensorRT-8.2.3.0") #7 set(TRT_INCLUDE_DIRS ${TRT_DIR}\\include) #8 set(TRT_LIB_DIRS ${TRT_DIR}\\lib) #9 set(Dirent_INCLUDE_DIRS "Z:\\code\\dirent-master\\include") #10add_definitions(-std=c++11) add_definitions(-DAPI_EXPORTS)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 QUIETNO_MODULENO_DEFAULT_PATHNO_CMAKE_PATHNO_CMAKE_ENVIRONMENT_PATHNO_SYSTEM_ENVIRONMENT_PATHNO_CMAKE_PACKAGE_REGISTRYNO_CMAKE_BUILDS_PATHNO_CMAKE_SYSTEM_PATHNO_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 ${PROJECT_SOURCE_DIR}/preprocess.cu ${PROJECT_SOURCE_DIR}/preprocess.h) #17target_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) if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) set(CMAKE_CUDA_ARCHITECTURES 70 75 80 86) endif(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
这里需要注意:
-
在tensorrtx下的yolov5文件夹中,打开cmd,输入以下代码
mkdir build cd build cmake ..
-
进入build文件夹,用vs打开yolov5.sln文件,然后点击生成,生成解决方案
出现以下这说明生成成功:
-
设置yolov5为启动项
-
进入yolov5属性页,调试选项设置如下:
-
设置完成运行程序
运行可能需要花费一定的时间,出现以下这说明模型转换成功
-
测试生成的模型,在属性页面中设置如下:
出现以下内容这说明运行成功: