锐客网

  1. 首页 > it技术 > >

win10系统下基于tensorrt的yolov5l网络部署

cuda


yolov5网络部署

  • 环境和使用项目
  • 模型训练
  • 生成wts文件
  • 构建生成vs项目
  • 修改CMakeList.txt
  • 生成engine文件
  • 用engine文件对图片进行推理
  • 存在的问题
  • 已解决的问题

环境和使用项目 环境：win10，vs2017，opencv3.4.14，cuda11.1.0，cudnn8.2.1，tensorrt7.2.3，pytorch1.9，安装tensorrt的方法见前文
yolov5 项目：https://github.com/ultralytics/yolov5
yolov5对应tensorrt项目tensorrtx：https://github.com/donnyyou/tensorrtx/tree/master/yolov5
模型训练 略
生成wts文件 【win10系统下基于tensorrt的yolov5l网络部署】根据yolov5训练出来的模型权重为yolov5l.pt，其他yolov5模型类似
根据pt文件生成wts文件：wts文件是从yolov5权重转换到tensorrx权重的中间文件。转换方法为：将yolov5l.pt文件和tensorrtx/yolov5下的gen_wts.py文件放在yolov5的根目录下（主要是为了使用yolov5项目下的相关python模块），python运行gen_wts.py文件，生成yolov5l.wts文件
构建生成vs项目 需要将tensorrtx项目构建成vs项目，使用cmake-gui构建生成项目，配置时注意config中的generator选择Visual Studio 15 2017，platform选择x64
修改CMakeList.txt tensorrtx项目是用于linux的C++项目，用于win10需要修改CMakeLists.txt文件，具体修改方法作者已给出。
cmake_minimum_required(VERSION 2.6)project(yolov5) # 1 set(OpenCV_DIR "D:\\opencv\\opencv346\\build")#2 set(TRT_DIR "D:\\TensorRT-7.0.0.11.Windows10.x86_64.cuda-10.2.cudnn7.6\\TensorRT-7.0.0.11")#3add_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) include_directories(${TRT_DIR}\\include)# -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 "libraries: ${OpenCV_LIBS}") message(STATUS "include path: ${OpenCV_INCLUDE_DIRS}")include_directories(${OpenCV_INCLUDE_DIRS}) link_directories(${TRT_DIR}\\lib)add_executable(yolov5 ${PROJECT_SOURCE_DIR}/yolov5.cpp ${PROJECT_SOURCE_DIR}/yololayer.cu ${PROJECT_SOURCE_DIR}/yololayer.h)#4target_link_libraries(yolov5 "nvinfer" "nvinfer_plugin")#5 target_link_libraries(yolov5 ${OpenCV_LIBS})#6 target_link_libraries(yolov5 ${CUDA_LIBRARIES})#7 target_link_libraries(yolov5 Threads::Threads)#8

注意重点是修改opencv和tensorrt的头文件和库文件目录
修改CMakeList.txt后，设置yolov5项目为启动项目
配置项目成功后，运行vs项目如图：
win10系统下基于tensorrt的yolov5l网络部署
文章图片

生成engine文件 engine文件是tensorrtx项目的权重文件，将yolov5l.wts放在vs项目下，添加启动项目时的命令行参数为：-s yolov5l.wts yolov5l.engine l，其中-s表示序列化模型对象，l表示yolov5l模型。重新运行项目即可获得yolov5l.engine文件。
win10系统下基于tensorrt的yolov5l网络部署
文章图片

这里实际生成engine文件花了40多分钟，是一个比较慢的速度。之前在jetson xavier平台上生成engine文件一般花2分钟左右的时间，可能和运行平台有关系。
用engine文件对图片进行推理 engine文件实际上保存的是c++用于推理的对象的权重数据，如果要对图像进行推理的话需要改动启动的命令行参数，改为：-d yolov5l.engine …/samples，其中-d表示反序列化文件，…/samples表示用于推理的图片的路径。运行后即可获得识别结果win10系统下基于tensorrt的yolov5l网络部署
文章图片

存在的问题 已解决的问题 win10平台下的tensorrtx工程生成engine时间过长，多等一下就好

    推荐阅读


    上一篇：#|动手学深度学习（3.14 正向传播、反向传播和计算图）

    下一篇：算法|PyTorch中的squeeze()和unsqueeze()详解与应用案例