一、部署环境
-
TensorRT 8.6
-
CUDA 11.7
-
TensorRTx
二、TensorRTx 核心优势
-
针对自定义网络做了比原生 TensorRT 更深度的优化。
-
在 NVIDIA 4090 上,YOLOv5s 的 TensorRTx 推理速度比原生 TensorRT 快 15-25%,FP16 精度下可达 350+ FPS(可参考该性能表现,YOLOv8 部署后性能优势一致)。
三、跨平台 / 跨版本兼容性
TensorRTx 封装了版本适配逻辑,同一套代码可兼容 TensorRT 8.0-8.6,且提供了跨平台的 CUDA 核函数实现,无需额外修改代码即可适配不同环境。
四、部署步骤
4.1 克隆 TensorRTx 仓库
bash
git clone https://github.com/wang-xinyu/tensorrtx.git
4.2 下载 YOLOv8 权重文件
bash
# 进入 YOLOv8 部署目录
cd tensorrtx/yolov8
# 下载 YOLOv8n.pt(轻量版,适合快速测试;也可替换为 s/m/l/x 版本)
wget https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt
4.3 将 .pt 权重转为 .wts 格式
bash
python gen_wts.py -w yolov8n.pt -o yolov8n.wts -t detect
4.4 配置参数并将 .wts 权重转为 .engine 格式
4.4.1 修改配置文件
打开路径~/yolov8_tensorrt/yolov8/include/config.h,根据自身需求修改以下参数(默认适配 COCO 数据集,无需额外修改可直接使用):
-
默认
kNumClass=80(对应 COCO 数据集 80 类别),若使用自定义数据集,需修改为自身数据集的类别数。 -
输入尺寸默认
kInputH=640、kInputW=640,匹配 YOLOv8 默认输入尺寸,无需修改。 -
默认量化方式为 FP16,适配 4090 性能最优,无需修改。
4.4.2 修改 CMakeLists.txt 并编译
bash
# 修改 CMakeLists.txt 中 TensorRT 路径(根据自身 TensorRT 安装路径调整)
# 编译构建
mkdir build
cd build
cp yolov8.wts build
cmake ..
make
4.4.3 链接库并执行转换
bash
# 链接 TensorRT 库(根据自身 TensorRT 安装路径调整)
export LD_LIBRARY_PATH=/usr/local/TensorRT-8.4.3.1/lib/:$LD_LIBRARY_PATH
# 执行权重转换(生成 yolov8.engine 文件)
./yolov8_det -s yolov8n.wts yolov8.engine n