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

總結

以上是生活随笔為你收集整理的YOLO-v5训练自己的数据+TensorRT推理部署（2）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。