YOLO-v5训练自己的数据+TensorRT推理部署

Posted 2022-03-07 踟蹰横渡口,彳亍上滩舟。

YOLO-v5训练自己的数据+TensorRT推理部署（2）

代码下载地址：下载地址

YOLO v5转TensorRT模型并调用

0.pt模型转wts模型

python3 gen_wts.py
# 注意修改代码中模型保存和模型加载的路径

1.修改部分文件

• 0.修改CMakeLists.txt

cmake_minimum_required(VERSION 2.6)
project(yolov5)
add_definitions(-std=c++11)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
find_package(CUDA REQUIRED)
set(CUDA_NVCC_PLAGS $CUDA_NVCC_PLAGS;-std=c++11;-g;-G;-gencode;arch=compute_30;code=sm_30)
include_directories($PROJECT_SOURCE_DIR/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib6