极智项目 | YOLO11目标检测算法训练+TensorRT部署实战

欢迎关注我的公众号 [极智视界]，获取我的更多技术分享

传送：极智项目 | YOLO11目标检测算法训练+TensorRT部署实战

大家好，我是极智视界，本文分享实战项目之YOLO11目标检测算法训练+TensorRT部署实战。

1. 项目介绍

项目作者：极智视界
项目init时间： 20241001
项目介绍：对YOLO11基于coco_minitrain_10k数据集进行训练，并使用py TensorRT进行加速推理，包括导onnx和onnx2trt转换脚本
项目参考：YOLO11部分参考 => https://github.com/ultralytics/ultralytics (这是YOLO11算法的官方出处)

2. 算法训练

(1) 数据集整备

数据集放在 datasets/coco_minitrain_10k 数据集目录结构如下：

datasets/
└── coco_mintrain_10k/
    ├── annotations/
    │   ├── instances_train2017.json
    │   ├── instances_val2017.json
    │   ├── ... (其他标注文件)
    ├── train2017/
    │   ├── 000000000001.jpg
    │   ├── ... (其他训练图像)
    ├── val2017/
    │   ├── 000000000001.jpg
    │   ├── ... (其他验证图像)
    └── test2017/
        ├── 000000000001.jpg
        ├── ... (其他测试图像)

(2) 训练环境搭建

conda creaet-n yolo11_py310 python=3.10

conda activate yolo11_py310

pip install -U -r train/requirements.txt

(3) 推理测试

先下载预训练权重：

执行预测测试：

预测结果保存在 runs 文件夹下，效果如下：

极智项目 | YOLO11目标检测算法训练+TensorRT部署实战_TensorRT

(4) 开启训练

已经准备好一键训练肩膀，直接执行训练脚本：

其中其作用的代码很简单，就在 train/train_yolo11.py 中，如下：

# Load a model
model = YOLO(curr_path + "/wgts/yolo11n.pt")

# Train the model
train_results = model.train(
    data= curr_path + "/cfg/coco128.yaml",  # path to dataset YAML
    epochs=100,  # number of training epochs
    imgsz=640,  # training image size
    device="0",  # device to run on, i.e. device=0 or device=0,1,2,3 or device=cpu
)

# Evaluate model performance on the validation set
metrics = model.val()

主要就是配置一下训练参数，如数据集路径、训练轮数、显卡ID、图片大小等，然后执行训练即可
训练完成后，训练日志会在 runs/train 文件夹下，比如训练中 val 预测图片如下：

极智项目 | YOLO11目标检测算法训练+TensorRT部署实战_TensorRT_02

这样就完成了算法训练

3. 算法部署

使用 TensorRT 进行算法部署

(1) 导ONNX

直接执行一键导出ONNX脚本:

在脚本中已经对ONNX做了sim的简化
生成的ONNX以及_simONNX模型保存在wgts文件夹下

(2) 安装tensorrt环境

直接去NVIDIA的官网下载( https://developer.nvidia.com/tensorrt/download)对应版本的tensorrt TAR包，解压
基本步骤如下：

tar zxvf TensorRT-xxx-.tar.gz

# 软链trtexec
sudo ln -s /path/to/TensorRT/bin/trtexec /usr/local/bin
# 验证一下
trtexec --help

# 安装trt的python接口
cd python
pip install tensorrt-xxx.whl

(3) 生成trt模型引擎文件

直接执行一键生成trt模型引擎的脚本：

正常会在wgts路径下生成yolo11n.engine，并有类似如下的日志：

[10/02/2024-21:28:48] [V] === Explanations of the performance metrics ===
[10/02/2024-21:28:48] [V] Total Host Walltime: the host walltime from when the first query (after warmups) is enqueued to when the last query is completed.
[10/02/2024-21:28:48] [V] GPU Compute Time: the GPU latency to execute the kernels for a query.
[10/02/2024-21:28:48] [V] Total GPU Compute Time: the summation of the GPU Compute Time of all the queries. If this is significantly shorter than Total Host Walltime, the GPU may be under-utilized because of host-side overheads or data transfers.
[10/02/2024-21:28:48] [V] Throughput: the observed throughput computed by dividing the number of queries by the Total Host Walltime. If this is significantly lower than the reciprocal of GPU Compute Time, the GPU may be under-utilized because of host-side overheads or data transfers.
[10/02/2024-21:28:48] [V] Enqueue Time: the host latency to enqueue a query. If this is longer than GPU Compute Time, the GPU may be under-utilized.
[10/02/2024-21:28:48] [V] H2D Latency: the latency for host-to-device data transfers for input tensors of a single query.
[10/02/2024-21:28:48] [V] D2H Latency: the latency for device-to-host data transfers for output tensors of a single query.
[10/02/2024-21:28:48] [V] Latency: the summation of H2D Latency, GPU Compute Time, and D2H Latency. This is the latency to infer a single query.
[10/02/2024-21:28:48] [I] 
&&&& PASSED TensorRT.trtexec [TensorRT v100500] [b18] # trtexec --onnx=../wgts/yolo11n_sim.onnx --saveEngine=../wgts/yolo11n.engine --fp16 --verbose

(4) 执行trt推理

直接执行一键推理脚本：

实际的trt推理脚本在 deploy/infer_trt.py 推理成功会有如下日志：

推理结果保存在 deploy/output.jpg

如下：

极智项目 | YOLO11目标检测算法训练+TensorRT部署实战_目标检测_03

好了，以上分享了实战项目之YOLO11目标检测算法训练+TensorRT部署实战，希望我的分享能对你的学习有一点帮助。